Using flag in place of thread context for internal search calls

Signed-off-by: Kshitij Tandon <[email protected]>

Author: tandonks

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupMetadataService.kt

@@ -242,7 +242,7 @@ class RollupMetadataService(
             val dateHistogram = rollup.dimensions.first() as DateHistogram
             val dateField = dateHistogram.sourceField
 
-            logger.info("Idhr se ja rha hu mai tumhe kya 2")
+            val bypassMarker = "_rollup_internal_bypass_${org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_SIZE_CHECK}"
 
             val searchRequest = SearchRequest(rollup.sourceIndex)
                 .source(
@@ -251,37 +251,31 @@ class RollupMetadataService(
                         .query(MatchAllQueryBuilder())
                         .sort("$dateField.date_histogram", SortOrder.ASC)
                         .trackTotalHits(false)
-                        .fetchSource(false)
+                        .fetchSource(org.opensearch.search.fetch.subphase.FetchSourceContext(false, arrayOf(bypassMarker), emptyArray()))
                         .docValueField("$dateField.date_histogram", DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT),
                 )
                 .allowPartialSearchResults(false)
 
-            // Set BYPASS_SIZE_CHECK to allow size=1 when querying rollup index to get earliest timestamp
-            // This is needed for continuous rollup initialization on rollup indices (multi-tier rollup)
-            org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.setBypass(
-                org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_SIZE_CHECK,
-            )
-            try {
-                val response: SearchResponse = client.suspendUntil { search(searchRequest, it) }
+            // Set bypass in ThreadContext for multi-node support - skips all rollup validations
+            // This is needed because we're querying a rollup index directly without aggregations
+            // Set bypass via preference parameter for multi-node support - allows size > 0 for fetching earliest timestamp
+//            searchRequest.preference("_rollup_bypass:${org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_SIZE_CHECK}")
 
-                if (response.hits.hits.isEmpty()) {
-                    return StartingTimeResult.NoDocumentsFound
-                }
+            val response: SearchResponse = client.suspendUntil { search(searchRequest, it) }
 
-                logger.info("Idhr se ja rha hu mai tumhe kya")
+            if (response.hits.hits.isEmpty()) {
+                return StartingTimeResult.NoDocumentsFound
+            }
 
-                // In rollup indices, date histogram fields are named as "field.date_histogram"
-                val rollupDateField = "$dateField.date_histogram"
-                val firstHitTimestampAsString: String =
-                    response.hits.hits.first().field(rollupDateField).getValue<String>()
-                        ?: return StartingTimeResult.NoDocumentsFound
+            // In rollup indices, date histogram fields are named as "field.date_histogram"
+            val rollupDateField = "$dateField.date_histogram"
+            val firstHitTimestampAsString: String =
+                response.hits.hits.first().field(rollupDateField).getValue<String>()
+                    ?: return StartingTimeResult.NoDocumentsFound
 
-                val formatter = DateFormatter.forPattern(DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT)
-                val epochMillis = DateFormatters.from(formatter.parse(firstHitTimestampAsString), formatter.locale()).toInstant().toEpochMilli()
-                return StartingTimeResult.Success(getRoundedTime(epochMillis, dateHistogram))
-            } finally {
-                org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.clearBypass()
-            }
+            val formatter = DateFormatter.forPattern(DATE_FIELD_STRICT_DATE_OPTIONAL_TIME_FORMAT)
+            val epochMillis = DateFormatters.from(formatter.parse(firstHitTimestampAsString), formatter.locale()).toInstant().toEpochMilli()
+            return StartingTimeResult.Success(getRoundedTime(epochMillis, dateHistogram))
         } catch (e: RemoteTransportException) {
             val unwrappedException = ExceptionsHelper.unwrapCause(e) as Exception
             logger.error("Error when getting earliest timestamp from rollup index for rollup [{}]: {}", rollup.id, unwrappedException)

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupRunner.kt

@@ -271,16 +271,7 @@ object RollupRunner :
                             withClosableContext(
                                 IndexManagementSecurityContext(job.id, settings, threadPool.threadContext, job.user),
                             ) {
-                                // Need to set this bypass as we are already doing the required aggregation re-writing in the composite aggregation
-                                // hence do not need to go through the interceptor for aggregation rewriting
-                                org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.setBypass(
-                                    org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.BYPASS_ROLLUP_SEARCH,
-                                )
-                                try {
-                                    rollupSearchService.executeCompositeSearch(updatableJob, metadata, clusterService)
-                                } finally {
-                                    org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor.clearBypass()
-                                }
+                                rollupSearchService.executeCompositeSearch(updatableJob, metadata, clusterService)
                             }
                         val rollupResult =
                             when (rollupSearchResult) {

src/main/kotlin/org/opensearch/indexmanagement/rollup/RollupSearchService.kt

@@ -19,13 +19,15 @@ import org.opensearch.core.action.ActionListener
 import org.opensearch.core.common.breaker.CircuitBreakingException
 import org.opensearch.indexmanagement.opensearchapi.retry
 import org.opensearch.indexmanagement.opensearchapi.suspendUntil
+import org.opensearch.indexmanagement.rollup.interceptor.RollupInterceptor
 import org.opensearch.indexmanagement.rollup.model.Rollup
 import org.opensearch.indexmanagement.rollup.model.RollupMetadata
 import org.opensearch.indexmanagement.rollup.settings.RollupSettings.Companion.MINIMUM_CANCEL_AFTER_TIME_INTERVAL_MINUTES
 import org.opensearch.indexmanagement.rollup.settings.RollupSettings.Companion.ROLLUP_SEARCH_BACKOFF_COUNT
 import org.opensearch.indexmanagement.rollup.settings.RollupSettings.Companion.ROLLUP_SEARCH_BACKOFF_MILLIS
 import org.opensearch.indexmanagement.rollup.util.getRollupSearchRequest
 import org.opensearch.search.aggregations.MultiBucketConsumerService
+import org.opensearch.search.fetch.subphase.FetchSourceContext
 import org.opensearch.transport.RemoteTransportException
 import org.opensearch.transport.client.Client
 import java.time.Instant
@@ -117,6 +119,14 @@ class RollupSearchService(
                     val cancelTimeoutTimeValue = TimeValue.timeValueMinutes(getCancelAfterTimeInterval(cancelAfterTimeInterval.minutes))
                     searchRequest.cancelAfterTimeInterval = cancelTimeoutTimeValue
 
+                    // Set bypass flag via FetchSourceContext for multi-node support
+                    // Using fetchSource=false with marker in includes to avoid affecting actual source fetching
+                    val bypassMarker = "_rollup_internal_bypass_${RollupInterceptor.BYPASS_ROLLUP_SEARCH}"
+
+                    searchRequest.source().fetchSource(
+                        FetchSourceContext(false, arrayOf(bypassMarker), emptyArray()),
+                    )
+
                     search(searchRequest, listener)
                 }
             },

src/main/kotlin/org/opensearch/indexmanagement/rollup/interceptor/RollupInterceptor.kt

@@ -69,6 +69,7 @@ class RollupInterceptor(
         /**
          * Thread-local bypass mechanism for internal operations that need to query rollup indices
          * without triggering interceptor validations.
+         * Kept for backward compatibility with single-node clusters.
          */
         private val bypassInterceptor = ThreadLocal<Int>()
 
@@ -108,6 +109,21 @@ class RollupInterceptor(
         fun getBypassLevel(): Int = bypassInterceptor.get() ?: 0
     }
 
+    /**
+     * Reads the bypass value from the request's FetchSourceContext.
+     * Returns the bypass level if the marker is present in includes array, null otherwise.
+     * This enables bypass mechanism to work in multi-node clusters.
+     */
+    private fun getBypassFromFetchSource(request: ShardSearchRequest): Int? {
+        val includes = request.source()?.fetchSource()?.includes()
+
+        // Look for our bypass marker in the includes array and extract the bypass level
+        return includes
+            ?.find { it.startsWith("_rollup_internal_bypass_") }
+            ?.substringAfter("_rollup_internal_bypass_")
+            ?.toIntOrNull()
+    }
+
     init {
         clusterService.clusterSettings.addSettingsUpdateConsumer(RollupSettings.ROLLUP_SEARCH_ENABLED) {
             searchEnabled = it
@@ -129,14 +145,26 @@ class RollupInterceptor(
     ): TransportRequestHandler<T> = object : TransportRequestHandler<T> {
         override fun messageReceived(request: T, channel: TransportChannel, task: Task) {
             if (searchEnabled && request is ShardSearchRequest) {
+                logger.info("Idhr aaya hu 3")
                 val index = request.shardId().indexName
                 val isRollupIndex = isRollupIndex(index, clusterService.state())
                 if (isRollupIndex) {
-                    val bypassLevel = getBypassLevel()
+                    // Check bypass from FetchSourceContext (multi-node) OR ThreadLocal (single-node/backward compat)
+                    val bypassFromFetchSource = getBypassFromFetchSource(request)
+                    val bypassFromThreadLocal = getBypassLevel()
+
+                    // Use FetchSourceContext value if present, otherwise fall back to ThreadLocal
+                    val effectiveBypass = bypassFromFetchSource ?: bypassFromThreadLocal
+
+                    logger.debug(
+                        "RollupInterceptor bypass check - fetchSource: $bypassFromFetchSource, " +
+                            "threadLocal: $bypassFromThreadLocal, effective: $effectiveBypass",
+                    )
+
                     // BYPASS_ROLLUP_SEARCH: Skip all validations and query rewriting
                     // Used for composite aggregation query which we do for retrieving the buckets
                     // to write to target index during rollup of rolled up index
-                    if (bypassLevel == BYPASS_ROLLUP_SEARCH) {
+                    if (effectiveBypass == BYPASS_ROLLUP_SEARCH) {
                         actualHandler.messageReceived(request, channel, task)
                         return
                     }
@@ -145,7 +173,7 @@ class RollupInterceptor(
                     // fetch documents (e.g., fetching earliest timestamp document for rolled index
                     // in case of continuous rollup job). Normal rollup searches must have size=0
                     // since they should only return aggregations
-                    if (bypassLevel != BYPASS_SIZE_CHECK && request.source().size() != 0) {
+                    if (effectiveBypass != BYPASS_SIZE_CHECK && request.source().size() != 0) {
                         throw IllegalArgumentException("Rollup search must have size explicitly set to 0, but found ${request.source().size()}")
                     }