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

@@ -48,11 +48,9 @@ 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
@@ -64,7 +62,6 @@ 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(
@@ -187,48 +184,6 @@ 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.
@@ -237,14 +192,6 @@ 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

@@ -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)
             }

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

@@ -39,10 +39,10 @@ import org.opensearch.indexmanagement.transform.model.ShardNewDocuments
 import org.opensearch.indexmanagement.transform.model.Transform
 import org.opensearch.indexmanagement.transform.model.TransformSearchResult
 import org.opensearch.indexmanagement.transform.model.TransformStats
+import org.opensearch.indexmanagement.transform.opensearchapi.retryTransformSearch
 import org.opensearch.indexmanagement.transform.settings.TransformSettings.Companion.TRANSFORM_JOB_SEARCH_BACKOFF_COUNT
 import org.opensearch.indexmanagement.transform.settings.TransformSettings.Companion.TRANSFORM_JOB_SEARCH_BACKOFF_MILLIS
 import org.opensearch.indexmanagement.transform.util.TransformContext
-import org.opensearch.indexmanagement.transform.util.TransformLockManager
 import org.opensearch.indexmanagement.util.IndexUtils.Companion.LUCENE_MAX_CLAUSES
 import org.opensearch.indexmanagement.util.IndexUtils.Companion.ODFE_MAGIC_NULL
 import org.opensearch.indexmanagement.util.IndexUtils.Companion.hashToFixedSize
@@ -61,6 +61,7 @@ import org.opensearch.search.aggregations.metrics.Percentiles
 import org.opensearch.search.aggregations.metrics.ScriptedMetric
 import org.opensearch.search.builder.SearchSourceBuilder
 import org.opensearch.transport.RemoteTransportException
+import java.time.Instant
 import java.util.concurrent.TimeUnit
 import kotlin.math.max
 import kotlin.math.pow
@@ -114,25 +115,39 @@ class TransformSearchService(
     }
 
     @Suppress("RethrowCaughtException")
-    suspend fun getShardLevelModifiedBuckets(transform: Transform, afterKey: Map<String, Any>?, currentShard: ShardNewDocuments, transformContext: TransformContext): BucketSearchResult {
+    suspend fun getShardLevelModifiedBuckets(
+        transform: Transform,
+        afterKey: Map<String, Any>?,
+        currentShard: ShardNewDocuments,
+        transformContext: TransformContext
+    ): BucketSearchResult {
         try {
             var retryAttempt = 0
             var pageSize = calculateMaxPageSize(transform)
-            val searchResponse = backoffPolicy.retry(logger, transformContext.transformLockManager) {
+            val searchStart = Instant.now().epochSecond
+            val searchResponse = backoffPolicy.retryTransformSearch(logger, transformContext.transformLockManager) {
                 val pageSizeDecay = 2f.pow(retryAttempt++)
+                val searchRequestTimeoutInSeconds = transformContext.getMaxRequestTimeoutInSeconds()
                 client.suspendUntil { listener: ActionListener<SearchResponse> ->
-                    pageSize = transformContext.pageSize ?: max(1, pageSize.div(pageSizeDecay.toInt()))
+                    // If the previous request of the current transform job execution was successful, take the page size of previous request.
+                    // If not, calculate the page size.
+                    pageSize = transformContext.lastSuccessfulPageSize ?: max(1, pageSize.div(pageSizeDecay.toInt()))
                     if (retryAttempt > 1) {
                         logger.debug(
                             "Attempt [${retryAttempt - 1}] to get modified buckets for transform [${transform.id}]. Attempting " +
                                 "again with reduced page size [$pageSize]"
                         )
                     }
-                    val request = getShardLevelBucketsSearchRequest(transform, afterKey, pageSize, currentShard, transformContext.getMaxRequestTimeout())
+                    if (searchRequestTimeoutInSeconds == null) {
+                        return@suspendUntil
+                    }
+                    val request = getShardLevelBucketsSearchRequest(transform, afterKey, pageSize, currentShard, searchRequestTimeoutInSeconds)
                     search(request, listener)
                 }
             }
-            transformContext.pageSize = pageSize
+            // If the request was successful, update page size
+            transformContext.lastSuccessfulPageSize = pageSize
+            transformContext.renewLockForLongSearch(Instant.now().epochSecond - searchStart)
             return convertBucketSearchResponse(transform, searchResponse)
         } catch (e: TransformSearchServiceException) {
             throw e
@@ -169,22 +184,28 @@ class TransformSearchService(
                     modifiedBuckets.size
 
             var retryAttempt = 0
-            val searchResponse = backoffPolicy.retry(logger, transformContext.transformLockManager) {
-                // TODO: Should we store the value of the past successful page size (?)
-                val pageSizeDecay =  2f.pow(retryAttempt++)
+            val searchStart = Instant.now().epochSecond
+            val searchResponse = backoffPolicy.retryTransformSearch(logger, transformContext.transformLockManager) {
+                val pageSizeDecay = 2f.pow(retryAttempt++)
+                val searchRequestTimeoutInSeconds = transformContext.getMaxRequestTimeoutInSeconds()
+
                 client.suspendUntil { listener: ActionListener<SearchResponse> ->
-                    pageSize = transformContext.pageSize ?: max(1, pageSize.div(pageSizeDecay.toInt()))
+                    // If the previous request of the current transform job execution was successful, take the page size of previous request.
+                    // If not, calculate the page size.
+                    pageSize = transformContext.lastSuccessfulPageSize ?: max(1, pageSize.div(pageSizeDecay.toInt()))
                     if (retryAttempt > 1) {
                         logger.debug(
                             "Attempt [${retryAttempt - 1}] of composite search failed for transform [${transform.id}]. Attempting " +
                                 "again with reduced page size [$pageSize]"
                         )
                     }
-                    val request = getSearchServiceRequest(transform, afterKey, pageSize, modifiedBuckets, transformContext.getMaxRequestTimeout())
+                    val request = getSearchServiceRequest(transform, afterKey, pageSize, modifiedBuckets, searchRequestTimeoutInSeconds)
                     search(request, listener)
                 }
             }
-            transformContext.pageSize = pageSize
+            // If the request was successful, update page size
+            transformContext.lastSuccessfulPageSize = pageSize
+            transformContext.renewLockForLongSearch(Instant.now().epochSecond - searchStart)
             return convertResponse(transform, searchResponse, modifiedBuckets = modifiedBuckets)
         } catch (e: TransformSearchServiceException) {
             throw e
@@ -210,7 +231,7 @@ class TransformSearchService(
             afterKey: Map<String, Any>? = null,
             pageSize: Int,
             modifiedBuckets: MutableSet<Map<String, Any>>? = null,
-            timeout: Long? = null
+            timeoutInSeconds: Long? = null
         ): SearchRequest {
             val sources = mutableListOf<CompositeValuesSourceBuilder<*>>()
             transform.groups.forEach { group -> sources.add(group.toSourceBuilder().missingBucket(true)) }
@@ -223,7 +244,7 @@ class TransformSearchService(
             } else {
                 getQueryWithModifiedBuckets(transform.dataSelectionQuery, modifiedBuckets, transform.groups)
             }
-            return getSearchServiceRequest(transform.sourceIndex, query, aggregationBuilder, timeout)
+            return getSearchServiceRequest(transform.sourceIndex, query, aggregationBuilder, timeoutInSeconds)
         }
 
         private fun getQueryWithModifiedBuckets(
@@ -251,7 +272,24 @@ class TransformSearchService(
             return query
         }
 
-        private fun getSearchServiceRequest(index: String, query: QueryBuilder, aggregationBuilder: CompositeAggregationBuilder, timeout: Long? = null): SearchRequest {
+        /**
+         * Creates transform search request and sets timeout if it is provided
+         * Referring on: https://github.com/opensearch-project/OpenSearch/pull/1085
+         * https://github.com/opensearch-project/documentation-website/blob/main/_opensearch/rest-api/search.md#url-parameters
+         * cancel_after_time_interval property is used in order to set timeout of transform search request has not been ported to version 1.0
+         * thus we can't use it for version 1.0 support
+         *
+         * @param index - index that will be searched
+         * @param query - any additional [RestStatus] values that should be retried
+         * @param aggregationBuilder - search aggregations
+         * @param timeoutInSeconds - timeout period used for transform search request
+         */
+        private fun getSearchServiceRequest(
+            index: String,
+            query: QueryBuilder,
+            aggregationBuilder: CompositeAggregationBuilder,
+            timeoutInSeconds: Long? = null
+        ): SearchRequest {
             val searchSourceBuilder = SearchSourceBuilder()
                 .trackTotalHits(false)
                 .size(0)
@@ -260,10 +298,9 @@ class TransformSearchService(
             val request = SearchRequest(index)
                 .source(searchSourceBuilder)
                 .allowPartialSearchResults(false)
-
-            if(timeout != null){
-                request.cancelAfterTimeInterval = TimeValue(timeout, TimeUnit.SECONDS)
-            }
+            // The time after which the search request will be canceled.
+            // Request-level parameter takes precedence over cancel_after_time_interval cluster setting. Default is -1.
+            request.cancelAfterTimeInterval = timeoutInSeconds?.let { TimeValue(timeoutInSeconds, TimeUnit.SECONDS) }
             return request
         }
 
@@ -272,15 +309,15 @@ class TransformSearchService(
             afterKey: Map<String, Any>? = null,
             pageSize: Int,
             currentShard: ShardNewDocuments,
-            timeout: Long?
+            timeoutInSeconds: Long?
         ): SearchRequest {
             val rangeQuery = getSeqNoRangeQuery(currentShard.from, currentShard.to)
             val query = QueryBuilders.boolQuery().filter(rangeQuery).must(transform.dataSelectionQuery)
             val sources = transform.groups.map { it.toSourceBuilder().missingBucket(true) }
             val aggregationBuilder = CompositeAggregationBuilder(transform.id, sources)
                 .size(pageSize)
                 .apply { afterKey?.let { this.aggregateAfter(it) } }
-            return getSearchServiceRequest(currentShard.shardId.indexName, query, aggregationBuilder, timeout)
+            return getSearchServiceRequest(currentShard.shardId.indexName, query, aggregationBuilder, timeoutInSeconds)
                 .preference("_shards:" + currentShard.shardId.id.toString())
         }
 

src/main/kotlin/org/opensearch/indexmanagement/transform/action/explain/TransportExplainTransformAction.kt

@@ -64,7 +64,7 @@ class TransportExplainTransformAction @Inject constructor(
 
     private val log = LogManager.getLogger(javaClass)
 
-    @Suppress("SpreadOperator", "NestedBlockDepth")
+    @Suppress("SpreadOperator", "NestedBlockDepth", "LongMethod")
     override fun doExecute(task: Task, request: ExplainTransformRequest, actionListener: ActionListener<ExplainTransformResponse>) {
         log.debug(
             "User and roles string from thread context: ${client.threadPool().threadContext.getTransient<String>(