Back to Intellij Community

Context Propagation

docs/notebooks/1-ContextPropagation.ipynb

2025.3-rc-217.8 KB
Original Source

Context Propagation

kotlin
%use intellij-platform
import com.intellij.util.application
import com.intellij.util.ui.EDT
import javax.swing.JOptionPane
import com.intellij.openapi.application.ModalityState
import com.intellij.openapi.progress.ProgressManager
import com.intellij.platform.ide.progress.ModalTaskOwner
import com.intellij.platform.ide.progress.runWithModalProgressBlocking
import com.intellij.openapi.progress.EmptyProgressIndicator
import com.intellij.openapi.progress.ProcessCanceledException
import java.util.concurrent.Callable
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.runBlocking
import kotlinx.coroutines.withContext
import kotlinx.coroutines.yield
import kotlin.coroutines.AbstractCoroutineContextElement
import kotlin.coroutines.CoroutineContext
import com.intellij.concurrency.currentThreadContext
import kotlin.coroutines.EmptyCoroutineContext
import com.intellij.concurrency.IntelliJContextElement
import com.intellij.openapi.progress.blockingContextScope
import com.intellij.openapi.progress.currentThreadCoroutineScope
import com.intellij.openapi.progress.withCurrentThreadCoroutineScopeBlocking
import kotlinx.coroutines.future.asCompletableFuture
import kotlinx.coroutines.CoroutineStart
import com.intellij.concurrency.installThreadContext
import kotlinx.coroutines.coroutineScope





In this notebook, we will present one of the core concepts in the IntelliJ Platform -- Context Propagation.

Motivation

As programming languages and frameworks get more high-level, they tend to operate with increasingly more abstract entities. It also concerns computations and control flow of programs. Java and Kotlin abstract machine code and OS-level threads into convenient syntax and objects associated with threads.

One of the convenient abstractions that programming languages provide is lexical scoping:

kotlin
fun sample(param: Int) {
  val x = param + 1
  if (2 + 2 == 4) {
    val y = x + 2
  }
  // `y` is not accessible here
  // but `param` is accessible
}

IntelliJ Platform is a multithreaded application, and it deals with computations rather than with single-threaded functions. The usual control flow (if / while / for) is extended with threads and asynchronous computations.

kotlin
application.executeOnPooledThread {
  application.assertIsNonDispatchThread() // we are on a background thread
  val intensiveComputation = 1 + 1
  application.invokeLater {
    application.assertIsDispatchThread() // we are on EDT
    JOptionPane.showMessageDialog(null, "Result of intensive computation: $intensiveComputation", "Sample title", JOptionPane.INFORMATION_MESSAGE)
  }
};

Sometimes in the Platform, we can notice a pattern: some data becomes associated with a computation, just like a variable is associated with its lexical scope.

Example: Modality State

Modality state is an object that indicates that a computation belongs to some modal window. The computation usually runs on background, and sometimes it may want to access the UI thread. In this case, we must pass modality state to the UI computation so that the platform will not postpone this execution.

kotlin
val outerModalityState = ModalityState.defaultModalityState()
// the default choice for modality is `nonModal`: every computation assumes that it is launched in absence of modal dialogs by default
assert(outerModalityState == ModalityState.nonModal())

// launching java-style modal progress
ProgressManager.getInstance().runProcessWithProgressSynchronously(
  {
    // modality is different here, as we entered a modal computation
    val innerModalityState = ModalityState.defaultModalityState()
    assert(innerModalityState != ModalityState.nonModal())

    Thread.sleep(1000) // emulation of long background work
    application.invokeAndWait({
                                val modalityStateInDispatchThread = ModalityState.defaultModalityState()
                                assert(innerModalityState == modalityStateInDispatchThread)
                              }, innerModalityState)
    Thread.sleep(1000) // another long background work
  }, "Sample modal dialog", true, null)

Exercise: what would happen if we passed ModalityState.nonModal() to invokeAndWait and why?

Example: (Legacy) Cancellation

Most of the computations in IntelliJ Platform need to be cancellable -- otherwise, the application becomes unresponsive. The legacy cancellation model in IntelliJ Platform is built around ProgressIndicator, which also can belong to several computations

kotlin
val indicator = EmptyProgressIndicator()

application.executeOnPooledThread(
  {
    Thread.sleep(1000)
    indicator.cancel()
  })

application.executeOnPooledThread(Callable<String> {
  return@Callable application.executeOnPooledThread(Callable<String> {
    Thread.sleep(500)
    ProgressManager.getInstance().runProcess(
      {
        try {
          while (true) {
            ProgressManager.checkCanceled()
          }
        }
        catch (e: ProcessCanceledException) {
        }
      }, indicator)
    return@Callable "canceled successfully"
  }).get()
}).get()

In both these examples we see an instance of a parameter that is attached to a computation, and this parameter should transcend the boundaries of a single thread. The idea behind Context Propagation is to provide a unified interface for such parameters.

Kotlin Coroutines

In Kotlin Coroutines, a similar problem was resolved with the help of CoroutineContext -- a heterogeneous list that is attached to computations and which is inherited after coroutine transitions.

kotlin
class MyCoroutineContextElement(val id: Int) : AbstractCoroutineContextElement(MyCoroutineContextElement) {
  companion object Key : CoroutineContext.Key<MyCoroutineContextElement>

  override fun toString(): String {
    return "MyCoroutineContextElement($id)"
  }
}

runBlocking {
  DISPLAY("Element in runBlocking: ${coroutineContext[MyCoroutineContextElement]}")
  withContext(MyCoroutineContextElement(1)) {
    DISPLAY("Element in withContext: ${coroutineContext[MyCoroutineContextElement]}")
    launch(Dispatchers.Default) {
      DISPLAY("Element in launch: ${coroutineContext[MyCoroutineContextElement]}")
      yield()
      DISPLAY("Element after yield: ${coroutineContext[MyCoroutineContextElement]}")
      withContext(MyCoroutineContextElement(2)) {
        DISPLAY("Element in another withContext: ${coroutineContext[MyCoroutineContextElement]}")
      }
    }
  }
}

Another instances of this pattern are Reader monad in Haskell and dynamic binding in Lisp-like languages.

Platform API

Context Propagation of IntelliJ Platform is a heterogeneous list of implicit parameters attached to a computation, alongside the rules of transferring of this list. Let's look at some examples.

Context Propagation is designed to be closely compatible with Kotlin coroutines. The list of parameters is implemented as CoroutineContext, and it can be retrieved with the function currentThreadContext

kotlin
DISPLAY(currentThreadContext())

The intended way to interact with context propagation is from suspending code. The coroutine context can be accessed with currentThreadContext() even in non-suspend functions.

kotlin
fun nonSuspendingFunction(action: () -> Unit) {
  action()
}

runBlocking {
  withContext(MyCoroutineContextElement(42)) {
    val elementFromCoroutines = coroutineContext[MyCoroutineContextElement]
    DISPLAY("Element from coroutines: ${elementFromCoroutines}")
    nonSuspendingFunction {
      val elementFromThreadContext = currentThreadContext()[MyCoroutineContextElement] // no suspension here!
      DISPLAY("Element from thread context: $elementFromThreadContext")
    }
  }
}

Internally, context propagation is implemented as a thread-local variable containing CoroutineContext. Note, that coroutines machinery takes precedence over this thread local:

kotlin
runBlocking {
  withContext(MyCoroutineContextElement(42)) {
    DISPLAY("Element after withContext: ${currentThreadContext()[MyCoroutineContextElement]}")
    launch(start = CoroutineStart.UNDISPATCHED, context = MyCoroutineContextElement(43)) { // note the undispatched start!
      DISPLAY("Element after undispatched launch: ${currentThreadContext()[MyCoroutineContextElement]}")
    }
  }
}

If needed, thread context can be overridden from blocking code with installThreadContext. Note, that installThreadContext will replace the thread-local value by default, instead of appending to it (like it is done in withContext). The reason is that installThreadContext is mostly used by the Platform, and it often needs complete replacement of the context.

kotlin
installThreadContext(MyCoroutineContextElement(42)) {
  DISPLAY("Explicitly installed thread context: ${currentThreadContext()[MyCoroutineContextElement]}")
}

Cross-thread propagation

By default, coroutine context is not transferred to another thread. This is because old concurrency model of IntelliJ Platform was following Java, where asynchronous computations were not really tracked by anyone. We expect that average coroutine context elements are not prepared to be propagated to "fire-and-forget" computations:

kotlin
runBlocking {
  withContext(MyCoroutineContextElement(42)) {
    DISPLAY("Element context inside withContext: ${currentThreadContext()[MyCoroutineContextElement]}")
    application.invokeLater {
      DISPLAY("Element context inside invokeLater: ${currentThreadContext()[MyCoroutineContextElement]}")
    }
    application.invokeLater {
      DISPLAY("Element context inside executeOnPooledThread: ${currentThreadContext()[MyTransferrableElement]}")
    }
  }
}

IntelliJ Platform provides a way for some elements to permit their transfer via computations such as invokeLater. The coroutine context element needs to extend IntelliJContextElement interface.

kotlin
class MyTransferrableElement(val id: Int) : IntelliJContextElement, AbstractCoroutineContextElement(MyTransferrableElement) {
  companion object Key : CoroutineContext.Key<MyTransferrableElement>

  // The element permits its transfer here
  override fun produceChildElement(parentContext: CoroutineContext, isStructured: Boolean): IntelliJContextElement? {
    return this
  }

  override fun toString(): String {
    return "MyTransferrableElement($id)"
  }
}

installThreadContext(MyTransferrableElement(42)) {
  DISPLAY("Element context inside installThreadContext: ${currentThreadContext()[MyTransferrableElement]}")
  application.invokeLater {
    DISPLAY("Element context inside invokeLater: ${currentThreadContext()[MyTransferrableElement]}")
  }
  application.executeOnPooledThread {
    DISPLAY("Element context inside executeOnPooledThread: ${currentThreadContext()[MyTransferrableElement]}")
  }
}

Exercise: try to create a context element that gets propagated only if MyTransferrableElement is present at the moment of asynchronous scheduling (i.e., the invocation of invokeLater)

In particular, it means that asynchronous computations (such that a lambda in invokeLater) carry their own scope with them. Sometimes there is a need to execute computations synchronously -- the most popular example is synchronous dispatch of events by the EDT event queue. In this case, the Platform resets the context before event dispatch -- otherwise, there would be an error. The resetting of the context can be achieved with resetThreadContext function.

kotlin
installThreadContext(MyCoroutineContextElement(42)) {
  DISPLAY("Element context inside installThreadContext: ${currentThreadContext()[MyCoroutineContextElement]}")
  resetThreadContext {
    DISPLAY("Element context inside resetThreadContext: ${currentThreadContext()[MyCoroutineContextElement]}")
  }
}

Use case: Cancellation Propagation

One important application of context propagation is cancellation propagation. This is a subset of Context Propagation, which is used for maintaining the lifetime descriptor of computations and to get a limited version of structured concurrency. It is explored better in the notebook about Cancellation Model of IntelliJ Platform.

Use case: blockingContextScope

Some plugins in IntelliJ Platform were written against the legacy concurrency model. Modern implementation of the Platform is interested in controlling the scope of computations running in old plugins. For example, we can have a chain of asynchronous computations:

kotlin
application.invokeLater {
  application.executeOnPooledThread {
    application.invokeLater {
      application.executeOnPooledThread {
        application.invokeAndWait {
          DISPLAY("I am deep in asynchronous stack!")
        }
      }
    }
  }
}

To track these computations, the platform introduced a special function which is called blockingContextScope. This function is a combination of coroutineScope and blockingContext, and it allows to track the legacy computations and await their termination:

kotlin
runBlocking {
  DISPLAY("Launching `blockingContextScope`")
  blockingContextScope {
    application.executeOnPooledThread {
      DISPLAY("First `executeOnPooledThread` is running")
      Thread.sleep(1000)
      DISPLAY("First `executeOnPooledThread` is finished")
    }
    DISPLAY("First `executeOnPooledThread` is scheduled")
  }
  DISPLAY("`blockingContextScope` is finished`")
}

blockingContextScope internally works via transferable elements.

Use case: Observation.awaitConfiguration

Some applications that are using IntelliJ Platform are interested in using a headless instance of an IDE to import and index a project, so that they later can query something about the codebase. IntelliJ Platform provides a headless runner that is called warmup, which opens a project and waits for all non-trivial processes are finished.

Internally, warmup enumerates a set of important entry points (i.e., computations that run on project open) and then waits for their completion. Some of these computations can execute something asynchronusly -- in this case, warmup also tracks these asynchrnous computations as well. This way warmup is aware about the tree of computations, and it considers the process of project configuration done when all computations in this tree terminate. This turns out to be a surprisingly stable way to await the configuration process of the IDE -- and it requires minimal external support. Warmup was one of the original motivations to implement Context Propagation.

The entry points are usually marked with com.intellij.platform.backend.observation.TrackingUtil.trackActivity, and it allows tracking the whole configuration process.

Advanced Topics

Structured Propagation

Sometimes we know that the child computation will not outlive the parent one. The simplest example is Application.invokeAndWait. If this is the case, then context propagation is structured.

kotlin
DISPLAY("Before invokeAndWait")
application.invokeAndWait {
  Thread.sleep(500)
  DISPLAY("Inside invokeAndWait")
  Thread.sleep(500)
}
DISPLAY("After invokeAndWait")

This case is similar to synchronous execution, and here all elements are propagated. Additionally, the inheritors of IntelliJContextElement can react to structured propagation via isStructured flag.

kotlin
class MyTransferrableStructuredElement(val id: Int) : IntelliJContextElement, AbstractCoroutineContextElement(MyTransferrableStructuredElement) {
  companion object Key : CoroutineContext.Key<MyTransferrableStructuredElement>

  // The element permits its transfer here
  override fun produceChildElement(parentContext: CoroutineContext, isStructured: Boolean): IntelliJContextElement? {
    if (isStructured) {
      return this
    }
    else {
      return null
    }
  }

  override fun toString(): String {
    return "MyTransferrableStructuredElement($id)"
  }
}

installThreadContext(MyTransferrableStructuredElement(42) + MyCoroutineContextElement(43)) {
  DISPLAY("Element context inside installThreadContext: Transferrable -- ${currentThreadContext()[MyTransferrableStructuredElement]}, Regular -- ${currentThreadContext()[MyCoroutineContextElement]}")
  application.invokeLater {
    DISPLAY("Element context inside invokeLater: Transferrable -- ${currentThreadContext()[MyTransferrableStructuredElement]}, Regular -- ${currentThreadContext()[MyCoroutineContextElement]}")
  }
  application.invokeAndWait {
    DISPLAY("Element context inside invokeAndWait: Transferrable -- ${currentThreadContext()[MyTransferrableStructuredElement]}, Regular -- ${currentThreadContext()[MyCoroutineContextElement]}")
  }
}

Context Checkpoints

In Kotlin Coroutines, there is a way to launch asynchronous computations that are bound to some CoroutineScope:

kotlin
runBlocking {
  coroutineScope {
    launch { DISPLAY("I am in a coroutine scope") }
    launch { DISPLAY("I too am in a coroutine scope") }
  }
}

It is difficult to use coroutine scope in non-suspending code, as it is unclear where the computations are bounded. However, the platform offers one solution here. It is possible to create a coroutine scope from the currently installed thread context with withCurrentThreadCoroutineScopeBlocking, and then access it with currentThreadCoroutineContext:

kotlin
installThreadContext(MyCoroutineContextElement(42)) {
  val (result, trackingJob) = withCurrentThreadCoroutineScopeBlocking {
    DISPLAY("Current context element: ${currentThreadContext()[MyCoroutineContextElement]}") // prints 42
    installThreadContext(currentThreadContext() + MyCoroutineContextElement(43), true) {
      DISPLAY("Context element after replacement: ${currentThreadContext()[MyCoroutineContextElement]}") // prints 43
      currentThreadCoroutineScope().launch {
        DISPLAY("Context inside `launch`: ${currentThreadContext()[MyCoroutineContextElement]}") // prints 42!
      }
    }
  }
  trackingJob.asCompletableFuture().join()
}

This is useful when you want to pass the ability to launch coroutines on some scope which is under your control