Replace usages of ThreadContext.stashContext with pluginSubject.runAs

build.gradle

@@ -91,6 +91,7 @@ plugins {
     id "de.undercouch.download" version "5.6.0"
     id "com.netflix.nebula.ospackage" version "12.0.0"
     id "com.dorongold.task-tree" version "2.1.1"
+    id "org.gradle.test-retry" version "1.6.2"
 }
 
 apply plugin: 'java'
@@ -146,6 +147,10 @@ opensearchplugin {
 
 tasks.named("integTest").configure {
     it.dependsOn(project.tasks.named("bundlePlugin"))
+    it.retry {
+        failOnPassedAfterRetry = false
+        maxRetries = 5
+    }
 }
 
 allOpen {

src/main/kotlin/org/opensearch/indexmanagement/IndexManagementPlugin.kt

@@ -29,6 +29,7 @@ import org.opensearch.core.xcontent.XContentParser.Token
 import org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken
 import org.opensearch.env.Environment
 import org.opensearch.env.NodeEnvironment
+import org.opensearch.identity.PluginSubject
 import org.opensearch.indexmanagement.IndexManagementIndices.Companion.HISTORY_ALL
 import org.opensearch.indexmanagement.controlcenter.notification.ControlCenterIndices
 import org.opensearch.indexmanagement.controlcenter.notification.action.delete.DeleteLRONConfigAction
@@ -175,13 +176,15 @@ import org.opensearch.indexmanagement.transform.resthandler.RestPreviewTransform
 import org.opensearch.indexmanagement.transform.resthandler.RestStartTransformAction
 import org.opensearch.indexmanagement.transform.resthandler.RestStopTransformAction
 import org.opensearch.indexmanagement.transform.settings.TransformSettings
+import org.opensearch.indexmanagement.util.PluginClient
 import org.opensearch.indices.SystemIndexDescriptor
 import org.opensearch.jobscheduler.spi.JobSchedulerExtension
 import org.opensearch.jobscheduler.spi.ScheduledJobParser
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner
 import org.opensearch.monitor.jvm.JvmService
 import org.opensearch.plugins.ActionPlugin
 import org.opensearch.plugins.ExtensiblePlugin
+import org.opensearch.plugins.IdentityAwarePlugin
 import org.opensearch.plugins.NetworkPlugin
 import org.opensearch.plugins.Plugin
 import org.opensearch.plugins.SystemIndexPlugin
@@ -208,7 +211,8 @@ class IndexManagementPlugin :
     ActionPlugin,
     ExtensiblePlugin,
     SystemIndexPlugin,
-    TelemetryAwarePlugin {
+    TelemetryAwarePlugin,
+    IdentityAwarePlugin {
     private val logger = LogManager.getLogger(javaClass)
     lateinit var indexManagementIndices: IndexManagementIndices
     lateinit var actionValidation: ActionValidation
@@ -223,6 +227,7 @@ class IndexManagementPlugin :
     private val extensionCheckerMap = mutableMapOf<String, StatusChecker>()
     lateinit var indexOperationActionFilter: IndexOperationActionFilter
     private lateinit var metricsRegistry: MetricsRegistry
+    private lateinit var pluginClient: PluginClient
 
     companion object {
         const val PLUGINS_BASE_URI = "/_plugins"
@@ -398,6 +403,8 @@ class IndexManagementPlugin :
             environment,
         )
 
+        this.pluginClient = PluginClient(client)
+
         IndexManagementActionsMetrics.instance.initialize(metricsRegistry)
         rollupInterceptor = RollupInterceptor(clusterService, settings, indexNameExpressionResolver)
         val jvmService = JvmService(environment.settings())
@@ -430,13 +437,13 @@ class IndexManagementPlugin :
                 .registerMetadataServices(RollupMetadataService(client, xContentRegistry))
                 .registerConsumers()
                 .registerClusterConfigurationProvider(skipFlag)
-        indexManagementIndices = IndexManagementIndices(settings, client.admin().indices(), clusterService)
-        val controlCenterIndices = ControlCenterIndices(client.admin().indices(), clusterService)
+        indexManagementIndices = IndexManagementIndices(settings, this.pluginClient.admin().indices(), clusterService)
+        val controlCenterIndices = ControlCenterIndices(this.pluginClient.admin().indices(), clusterService)
         actionValidation = ActionValidation(settings, clusterService, jvmService)
         val indexStateManagementHistory =
             IndexStateManagementHistory(
                 settings,
-                client,
+                this.pluginClient,
                 threadPool,
                 clusterService,
                 indexManagementIndices,
@@ -454,7 +461,7 @@ class IndexManagementPlugin :
         val extensionChecker = ExtensionStatusChecker(extensionCheckerMap, clusterService)
         val managedIndexRunner =
             ManagedIndexRunner
-                .registerClient(client)
+                .registerClient(pluginClient)
                 .registerClusterService(clusterService)
                 .registerValidationService(actionValidation)
                 .registerNamedXContentRegistry(xContentRegistry)
@@ -472,7 +479,7 @@ class IndexManagementPlugin :
         val managedIndexCoordinator =
             ManagedIndexCoordinator(
                 environment.settings(),
-                client, clusterService, threadPool, indexManagementIndices, indexMetadataProvider, xContentRegistry,
+                pluginClient, clusterService, threadPool, indexManagementIndices, indexMetadataProvider, xContentRegistry,
             )
 
         val smRunner = SMRunner.init(client, threadPool, settings, indexManagementIndices, clusterService)
@@ -500,6 +507,7 @@ class IndexManagementPlugin :
             indexMetadataProvider,
             smRunner,
             pluginVersionSweepCoordinator,
+            pluginClient,
         )
     }
 
@@ -614,6 +622,10 @@ class IndexManagementPlugin :
         ActionPlugin.ActionHandler(DeleteLRONConfigAction.INSTANCE, TransportDeleteLRONConfigAction::class.java),
     )
 
+    override fun assignSubject(pluginSubject: PluginSubject?) {
+        pluginClient.setSubject(pluginSubject)
+    }
+
     override fun getTransportInterceptors(
         namedWriteableRegistry: NamedWriteableRegistry,
         threadContext: ThreadContext,

src/main/kotlin/org/opensearch/indexmanagement/common/model/notification/Channel.kt

@@ -79,20 +79,21 @@ data class Channel(val id: String) :
         user: User?,
     ) {
         val channel = this
-        client.threadPool().threadContext.stashContext().use {
-            // We need to set the user context information in the thread context for notification plugin to correctly resolve the user object
+        // TODO Understand why this is called twice when reindexing is finished in NotificationActionListenerIT.test notify for reindex with runtime policy
+        // We need to set the user context information in the thread context for notification plugin to correctly resolve the user object
+        client.threadPool().threadContext.getTransient<String>(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT)?.let {
             client.threadPool().threadContext.putTransient(ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT, generateUserString(user))
-            val res: SendNotificationResponse =
-                NotificationsPluginInterface.suspendUntil {
-                    this.sendNotification(
-                        (client as NodeClient),
-                        eventSource,
-                        ChannelMessage(message, null, null),
-                        listOf(channel.id),
-                        it,
-                    )
-                }
-            validateResponseStatus(res.getStatus(), res.notificationEvent.eventSource.referenceId)
         }
+        val res: SendNotificationResponse =
+            NotificationsPluginInterface.suspendUntil {
+                this.sendNotification(
+                    (client as NodeClient),
+                    eventSource,
+                    ChannelMessage(message, null, null),
+                    listOf(channel.id),
+                    it,
+                )
+            }
+        validateResponseStatus(res.getStatus(), res.notificationEvent.eventSource.referenceId)
     }
 }

src/main/kotlin/org/opensearch/indexmanagement/controlcenter/notification/action/delete/TransportDeleteLRONConfigAction.kt

@@ -15,16 +15,17 @@ import org.opensearch.common.inject.Inject
 import org.opensearch.commons.ConfigConstants
 import org.opensearch.core.action.ActionListener
 import org.opensearch.indexmanagement.IndexManagementPlugin
+import org.opensearch.indexmanagement.util.PluginClient
 import org.opensearch.tasks.Task
 import org.opensearch.transport.TransportService
-import org.opensearch.transport.client.node.NodeClient
+import org.opensearch.transport.client.Client
 
 class TransportDeleteLRONConfigAction
 @Inject
 constructor(
-    val client: NodeClient,
     transportService: TransportService,
     actionFilters: ActionFilters,
+    val client: PluginClient,
 ) : HandledTransportAction<DeleteLRONConfigRequest, DeleteResponse>(
     DeleteLRONConfigAction.NAME, transportService, actionFilters, ::DeleteLRONConfigRequest,
 ) {
@@ -35,7 +36,7 @@ constructor(
     }
 
     inner class DeleteLRONConfigHandler(
-        private val client: NodeClient,
+        private val client: Client,
         private val actionListener: ActionListener<DeleteResponse>,
         private val request: DeleteLRONConfigRequest,
         private val docId: String = request.docId,
@@ -49,13 +50,11 @@ constructor(
                 }",
             )
 
-            client.threadPool().threadContext.stashContext().use {
-                val deleteRequest =
-                    DeleteRequest(IndexManagementPlugin.CONTROL_CENTER_INDEX, docId)
-                        .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+            val deleteRequest =
+                DeleteRequest(IndexManagementPlugin.CONTROL_CENTER_INDEX, docId)
+                    .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
 
-                client.delete(deleteRequest, actionListener)
-            }
+            client.delete(deleteRequest, actionListener)
         }
     }
 }

src/main/kotlin/org/opensearch/indexmanagement/controlcenter/notification/action/get/TransportGetLRONConfigAction.kt

@@ -26,15 +26,16 @@ import org.opensearch.indexmanagement.controlcenter.notification.LRONConfigRespo
 import org.opensearch.indexmanagement.controlcenter.notification.model.LRONConfig
 import org.opensearch.indexmanagement.controlcenter.notification.util.getLRONConfigAndParse
 import org.opensearch.indexmanagement.opensearchapi.parseWithType
+import org.opensearch.indexmanagement.util.PluginClient
 import org.opensearch.search.builder.SearchSourceBuilder
 import org.opensearch.tasks.Task
 import org.opensearch.transport.TransportService
-import org.opensearch.transport.client.node.NodeClient
+import org.opensearch.transport.client.Client
 
 class TransportGetLRONConfigAction
 @Inject
 constructor(
-    val client: NodeClient,
+    val client: PluginClient,
     transportService: TransportService,
     actionFilters: ActionFilters,
     val xContentRegistry: NamedXContentRegistry,
@@ -48,7 +49,7 @@ constructor(
     }
 
     inner class GetLRONConfigHandler(
-        private val client: NodeClient,
+        private val client: Client,
         private val actionListener: ActionListener<GetLRONConfigResponse>,
         private val request: GetLRONConfigRequest,
     ) {
@@ -58,25 +59,23 @@ constructor(
                     ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT,
                 )}",
             )
-            client.threadPool().threadContext.stashContext().use {
-                if (null != request.docId) {
-                    getLRONConfigAndParse(
-                        client,
-                        request.docId,
-                        xContentRegistry,
-                        object : ActionListener<LRONConfigResponse> {
-                            override fun onResponse(response: LRONConfigResponse) {
-                                actionListener.onResponse(GetLRONConfigResponse(listOf(response), 1))
-                            }
+            if (null != request.docId) {
+                getLRONConfigAndParse(
+                    client,
+                    request.docId,
+                    xContentRegistry,
+                    object : ActionListener<LRONConfigResponse> {
+                        override fun onResponse(response: LRONConfigResponse) {
+                            actionListener.onResponse(GetLRONConfigResponse(listOf(response), 1))
+                        }
 
-                            override fun onFailure(e: Exception) {
-                                actionListener.onFailure(ExceptionsHelper.unwrapCause(e) as Exception)
-                            }
-                        },
-                    )
-                } else {
-                    doSearch()
-                }
+                        override fun onFailure(e: Exception) {
+                            actionListener.onFailure(ExceptionsHelper.unwrapCause(e) as Exception)
+                        }
+                    },
+                )
+            } else {
+                doSearch()
             }
         }
 

src/main/kotlin/org/opensearch/indexmanagement/controlcenter/notification/action/index/TransportIndexLRONConfigAction.kt

@@ -28,16 +28,17 @@ import org.opensearch.indexmanagement.controlcenter.notification.ControlCenterIn
 import org.opensearch.indexmanagement.controlcenter.notification.LRONConfigResponse
 import org.opensearch.indexmanagement.controlcenter.notification.util.getDocID
 import org.opensearch.indexmanagement.controlcenter.notification.util.getPriority
+import org.opensearch.indexmanagement.util.PluginClient
 import org.opensearch.indexmanagement.util.SecurityUtils
 import org.opensearch.tasks.Task
 import org.opensearch.transport.TransportService
-import org.opensearch.transport.client.node.NodeClient
+import org.opensearch.transport.client.Client
 
 @Suppress("LongParameterList")
 class TransportIndexLRONConfigAction
 @Inject
 constructor(
-    val client: NodeClient,
+    val client: PluginClient,
     transportService: TransportService,
     actionFilters: ActionFilters,
     val clusterService: ClusterService,
@@ -53,7 +54,7 @@ constructor(
     }
 
     inner class IndexLRONConfigHandler(
-        private val client: NodeClient,
+        private val client: Client,
         private val actionListener: ActionListener<LRONConfigResponse>,
         private val request: IndexLRONConfigRequest,
         private val user: User? = SecurityUtils.buildUser(client.threadPool().threadContext),
@@ -65,16 +66,14 @@ constructor(
                     ConfigConstants.OPENSEARCH_SECURITY_USER_INFO_THREAD_CONTEXT,
                 )}",
             )
-            client.threadPool().threadContext.stashContext().use {
-                // we use dryRun to help check permission and do request validation
-                if (request.dryRun) {
-                    validate()
-                    return
-                }
-                controlCenterIndices.checkAndUpdateControlCenterIndex(
-                    ActionListener.wrap(::onCreateMappingsResponse, actionListener::onFailure),
-                )
+            // we use dryRun to help check permission and do request validation
+            if (request.dryRun) {
+                validate()
+                return
             }
+            controlCenterIndices.checkAndUpdateControlCenterIndex(
+                ActionListener.wrap(::onCreateMappingsResponse, actionListener::onFailure),
+            )
         }
 
         private fun onCreateMappingsResponse(response: AcknowledgedResponse) {

src/main/kotlin/org/opensearch/indexmanagement/controlcenter/notification/filter/NotificationActionListener.kt

@@ -168,31 +168,30 @@ class NotificationActionListener<Request : ActionRequest, Response : ActionRespo
                 DEFAULT_PAGINATION_SIZE, 0, DEFAULT_LRON_CONFIG_SORT_FIELD, SORT_ORDER_DESC, queryString,
             )
 
-        client.threadPool().threadContext.stashContext().use {
-            client.execute(
-                GetLRONConfigAction.INSTANCE,
-                GetLRONConfigRequest(searchParams = searchParam),
-                object : ActionListener<GetLRONConfigResponse> {
-                    override fun onResponse(lronConfigResponse: GetLRONConfigResponse) {
-                        launch {
-                            sendNotification(lronConfigResponse, taskId, action, result)
-                        }
+        // TODO verify this works
+        client.execute(
+            GetLRONConfigAction.INSTANCE,
+            GetLRONConfigRequest(searchParams = searchParam),
+            object : ActionListener<GetLRONConfigResponse> {
+                override fun onResponse(lronConfigResponse: GetLRONConfigResponse) {
+                    launch {
+                        sendNotification(lronConfigResponse, taskId, action, result)
                     }
+                }
 
-                    override fun onFailure(e: Exception) {
-                        if (e is IndexNotFoundException) {
-                            logger.debug(
-                                "No notification policy configured for task: {} action: {}",
-                                taskId.toString(),
-                                action,
-                            )
-                        } else {
-                            logger.error("Can't get notification policy for action: {}", action, e)
-                        }
+                override fun onFailure(e: Exception) {
+                    if (e is IndexNotFoundException) {
+                        logger.debug(
+                            "No notification policy configured for task: {} action: {}",
+                            taskId.toString(),
+                            action,
+                        )
+                    } else {
+                        logger.error("Can't get notification policy for action: {}", action, e)
                     }
-                },
-            )
-        }
+                }
+            },
+        )
     }
 
     @Suppress("NestedBlockDepth")

src/main/kotlin/org/opensearch/indexmanagement/controlcenter/notification/util/LRONUtils.kt

@@ -22,7 +22,7 @@ import org.opensearch.indexmanagement.IndexManagementPlugin
 import org.opensearch.indexmanagement.controlcenter.notification.LRONConfigResponse
 import org.opensearch.indexmanagement.controlcenter.notification.model.LRONConfig
 import org.opensearch.indexmanagement.opensearchapi.parseFromGetResponse
-import org.opensearch.transport.client.node.NodeClient
+import org.opensearch.transport.client.Client
 
 const val LRON_DOC_ID_PREFIX = "LRON:"
 
@@ -68,7 +68,7 @@ fun getDocID(taskId: TaskId? = null, actionName: String? = null): String {
 }
 
 fun getLRONConfigAndParse(
-    client: NodeClient,
+    client: Client,
     docId: String,
     xContentRegistry: NamedXContentRegistry,
     actionListener: ActionListener<LRONConfigResponse>,