Optmize ContextN putAllInto

[koo-taejin]

Hello.

Motivation

I often use Coroutines and Reactors together.

To use these two together
I often use the Await.kt from https://github.com/Kotlin/kotlinx.coroutines
Unfortunately, there are performance issues.

Extract the ReactorContext from Await.kt.
Publisher has the extracted ReactorContext, and Subscriber has an empty ReactorContext because it creates a new one.
And finally combine these two ReactorContext in reactor.util.context.ContextN#putAllInto.

In my case, I always have 20 to 30 ReactorContexts.

In this case, base is empty, so the code below runs.

Context[] holder = new Context[]{base};
forEach((k, v) -> holder[0] = holder[0].put(k, v));
return holder[0];

If this code assumes that there are 20 contexts like me, many Contexts will be created.
In addition, because the previous value is copied again

1(Context1) (copy previous value 0)
2(Context2) (copy previous value  1)
3(Context3) (copy previous value  2)
4(Context4) (copy previous value  3)
5(Context5) (copy previous value  4)
6(ContextN) (copy previous value  5)
...
20 

To create ContextN, my application has to copy a lot.

Desired solution

• as-is

    Context[] holder = new Context[]{base};
    forEach((k, v) -> holder[0] = holder[0].put(k, v)); // Context is created as many times and copy previous values.
    return holder[0];

• desired solution

    final ContextN newContext = new ContextN(base.size() + this.size()); // ContextN object is created only once.
    newContext.putAll((Map<Object, Object>) this);
    base.forEach((k, v) -> newContext.accept(k, v));
    return newContext;

Considered alternatives

Additional context

PR : #3445

If you change this part like the changed code, resource usage can be reduced.
Anyway, since this is already ContextN, the result will be made into ContextN.
Only the number of Contexts will be copied, and Context objects are created only once, so it is judged that it will be possible to save resources much more than before.

Below is the code comparing when base is ContextN or not.
When running with AsyncProfiler individual tests on my MACBook

CPU

• as-is: 1,483,000
• to-be: 186,000
  Memory Allocation
• as-is: 28g
• to-be : 13g

class AwaitTest {

    val count = 10000000;

    @org.junit.Test
    fun testMonoAwaitContextPropagation() {
        val mapOf = mapOf(1 to "a", 2 to "b", 3 to "c", 4 to "d", 5 to "e", 6 to "f", 7 to "g")

        runBlocking(Context.of(mapOf).asCoroutineContext()) {
            val createMono = createMono()

            for (i in 1..count) {
                val injectCoroutineContext = injectCoroutineContext(createMono, coroutineContext)
                injectCoroutineContext.subscribe { v ->  }
            }
        }
    }


    @org.junit.Test
    fun testMonoAwaitContextPropagation2() {
        val mapOf = mapOf(1 to "a", 2 to "b", 3 to "c", 4 to "d", 5 to "e", 6 to "f", 7 to "g")

        runBlocking(Context.of(mapOf).asCoroutineContext()) {
            val createMono = createMono()

            for (i in 1..count) {
                val injectCoroutineContext = injectCoroutineContext2(createMono, coroutineContext)
                injectCoroutineContext.subscribe { v ->  }
            }
        }
    }


    private fun createMono(): Mono<String> = mono {
        val ctx = reactorContext()
        ctx.getOrDefault(7, "noValue")
    }

    private suspend fun reactorContext() = coroutineContext[ReactorContext]!!.context

    fun <T> injectCoroutineContext(publisher: Mono<T>, coroutineContext: CoroutineContext): Mono<T> {
        val reactorContext = coroutineContext[ReactorContext]?.context ?: return publisher
        return when (publisher) {
            is Mono -> publisher.contextWrite(reactorContext)
            else -> publisher
        }
    }

    fun <T> injectCoroutineContext2(publisher: Mono<T>, coroutineContext: CoroutineContext): Mono<T> {
        val reactorContext = coroutineContext[ReactorContext]?.context ?: return publisher
        return when (publisher) {
            is Mono -> publisher.contextWrite { c -> reactorContext.putAll(c) }
            else -> publisher
        }
    }
    
}

Thanks :)

[OlegDokuka]

fixed with #3448