Attach parent PluginResolver to plugin tracked coroutines (#4641)

Author: breedloj

plugins/core/jetbrains-community/src/migration/software/aws/toolkits/jetbrains/core/coroutines/PluginCoroutineScopeTracker.kt

@@ -10,9 +10,12 @@ import com.intellij.openapi.util.Disposer
 import kotlinx.coroutines.CoroutineName
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.ThreadContextElement
 import kotlinx.coroutines.cancel
 import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
+import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
 import java.util.concurrent.CancellationException
+import kotlin.coroutines.CoroutineContext
 
 class PluginCoroutineScopeTracker : Disposable {
     @PublishedApi
@@ -27,11 +30,30 @@ class PluginCoroutineScopeTracker : Disposable {
 }
 
 private class BackgroundThreadPoolScope(coroutineName: String, disposable: Disposable) : CoroutineScope {
-    override val coroutineContext = SupervisorJob() + CoroutineName(coroutineName) + getCoroutineBgContext()
+    override val coroutineContext = SupervisorJob() +
+        CoroutineName(coroutineName) +
+        getCoroutineBgContext() +
+        PluginResolverThreadContextElement(PluginResolver.fromCurrentThread())
 
     init {
         Disposer.register(disposable) {
             coroutineContext.cancel(CancellationException("Parent disposable was disposed"))
         }
     }
 }
+
+private class PluginResolverThreadContextElement(val pluginResolver: PluginResolver) : ThreadContextElement<PluginResolver> {
+    companion object Key : CoroutineContext.Key<PluginResolverThreadContextElement>
+
+    override val key = Key
+
+    override fun updateThreadContext(context: CoroutineContext): PluginResolver {
+        val oldPluginResolver = PluginResolver.fromCurrentThread()
+        PluginResolver.setThreadLocal(pluginResolver)
+        return oldPluginResolver
+    }
+
+    override fun restoreThreadContext(context: CoroutineContext, oldState: PluginResolver) {
+        PluginResolver.setThreadLocal(oldState)
+    }
+}

plugins/core/jetbrains-community/tst/software/aws/toolkits/jetbrains/core/coroutines/ScopeTest.kt

@@ -28,6 +28,7 @@ import org.junit.Rule
 import org.junit.Test
 import org.junit.rules.TemporaryFolder
 import org.junit.rules.TestName
+import software.aws.toolkits.jetbrains.services.telemetry.PluginResolver
 import software.aws.toolkits.jetbrains.utils.isInstanceOf
 import java.time.Duration
 import java.util.concurrent.CancellationException
@@ -153,6 +154,18 @@ class ScopeTest {
         assertThatThrownBy { disposableCoroutineScope(ApplicationManager.getApplication()) }.isInstanceOf<IllegalStateException>()
     }
 
+    @Test
+    fun `coroutine inherits PluginResolver from parent`() {
+        val pluginResolver = PluginResolver.fromCurrentThread()
+        PluginResolver.setThreadLocal(pluginResolver)
+
+        runBlocking {
+            withContext(applicationCoroutineScope().coroutineContext) {
+                assertThat(PluginResolver.fromCurrentThread()).isEqualTo(pluginResolver)
+            }
+        }
+    }
+
     private fun createFakePluginScope(componentManager: ComponentManager = ApplicationManager.getApplication()): Disposable {
         // We can't unload the real plugin in tests, so make another instance of the service and replace it for the tests
         val tracker = PluginCoroutineScopeTracker()