how to keep transaction when using `withContext` to setup a new security context?

Question

In the following demo we create a new ReactiveSecurityContextHolder to run an operation as a different user (in our actual code this is a bit more complex how we create the authentication, but for the demo the following shows the same problem as we experience).

@Service
final class SomeServiceImpl(
  private val repo: SomeRepository,
) : SomeService {

  @Transactional
  suspend fun foo() {
    val a = repo.save(Some())

    // expected: 1
    println(repo.count()); // actual: 1

    withContext(
      ReactorContext(
       ReactiveSecurityContextHolder.withAuthentication(
          RunAsUserToken(
            it.userId.toString(),
            it.userId,
            it.userId,
            mutableListOf(),
            AbstractAuthenticationToken::class.java,
          ),
        ),
      )
    ) {
      // expected: 1
      println(repo.count()); // actual: 0
    }
  }
}

Why does the the repo.count() within the withContext block not return 1 but 0?

When using withContext(Dispatchers.IO) the inner one returns 1 as well.

To me it seems like the transaction is lost during the context change, but I don't know how to keep it.

Answer 1

Answering my own question

TL;DR

This:

withContext(
      ReactorContext(...)
)

... overrides the ReactorContext key in the coroutineContext (note the lower-case c, coroutineContext refers to the CoroutineContext in the current coroutine). Keys are overriden by withContext and are not merged. The ReactorContext key will not be merged with an already existing ReactorContext but overridden. Hence the existing ReactorContext that holds the current transaction is hidden.

More explanation and solution

It is important to understand, that there is a ReactorContext in the coroutineContext and a new ReactorContext is created as argument to the withContext. These two ReactorContext exist independently of each other and require a manual merge to get a single ReactorContext that then can be merged to the coroutineContext (read: the existing one is override with the one that merges the items in the different ReactorContexts).

Here is an implementation that fixes the issue from the question:

@Service
final class SomeServiceImpl(
  private val repo: SomeRepository,
) : SomeService {

  @Transactional
  suspend fun foo() {
    val a = repo.save(Some())

    // expected: 1
    println(repo.count()); // actual: 1

    // get existing ReactorContext or create an empty one
    val reactorContext = (coroutineContext[ReactorContext] ?: ReactorContext(Context.empty()))

    // override security context in the reactorContext
   val reactorWithAuthContext = reactorContext.context.putAll(
      ReactiveSecurityContextHolder.withAuthentication(
          RunAsUserToken(
            it.userId.toString(),
            it.userId,
            it.userId,
            mutableListOf(),
            AbstractAuthenticationToken::class.java,
          ),
        ).readOnly(),
    )

    // merge the coroutineContext with the new ReactorContext which will hide the existing ReactorContext for the withContext block
    withContext(
     reactorWithAuthContext
    ) {
      // expected: 1
      println(repo.count()); // actual: 1
    }
  }
} 

The tricky parts for me to understand were:

1. ReactorContext is both, the name of the class but also (in kotlin) the name of the class refers to its companion object, which happens to be a Key that then is used to lookup something in the CouroutineContext, more clearly coroutineContext[ReactorContext] actually behaves like coroutineContext[ReactorContext.Key].

2. Neither are there different keys for each ReactorContext nor do the keys in the ReactorContext propagate to the CoroutineContext. The ReactorContext just wraps the reactor context and makes it available to the coroutine via the ReactorContext.Key in the CoroutineContext. I assumed that the mappings in a ReactorContext would automatically be merged into the CoroutineContext which is not true.