Context not propagated to onCompletion / onTermination / onFailure handlers

[ivu-mawi]

I'm trying to use Context passing to keep a transaction object across a Mutiny pipeline. It works great even across several transformToMulti calls and back to Uni etc.: the context is available. But when trying to use the context in an onCompletion or onTermination handler, I only get an empty context. onFailure works for Uni, but not Multi. I have not yet tested onCancellation.

I've tried to make some simple examples.
They're all the same except for how commitTransaction (or rollbackTransaction (identical for sake of example)) is invoked.

I would have hoped that in all examples, the same context (including the "transaction" object) is available inside the commitTransaction/rollbackTransaction callback. However, only onItem works in both Uni and Multi, and onFailure works only in Uni.

I'm guessing that onCompletion etc. make new internal pipelines with new subscribers where the context is not passed automatically.
Would this be a valid bug or feature request or am I missing something?

Note regarding alternatives to Mutiny context passing: I cannot use @RequestScoped because my pipelines are triggered by something that creates no request scope and I cannot use @Transactional because my database engine is not supported.

    public static void somethingEquivalentToMain() {
        contextAvailable_MultiOnItem().onFailure().recoverWithNull().await().indefinitely();
        contextMissing_MultiOnCompletion().onFailure().recoverWithNull().await().indefinitely();
        contextMissing_MultiOnFailure().onFailure().recoverWithNull().await().indefinitely();
        contextAvailable_UniOnItem().onFailure().recoverWithNull().await().indefinitely();
        contextMissing_UniOnTermination().onFailure().recoverWithNull().await().indefinitely();
        contextAvailable_UniOnFailure().onFailure().recoverWithNull().await().indefinitely();
    }

    public Uni<List<String>> contextAvailable_MultiOnItem() {
        log("\ncontextAvailable_MultiOnItem");
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // run some database queries
                .collect().asList()
                // commit transaction if there was no error
                .toMulti()
                .onItem().call(this::commitTransaction)
                .toUni()
                // some other async step
                .call(items -> Uni.createFrom().item(items))
                .log();
    }

    public Uni<List<String>> contextMissing_MultiOnCompletion() {
        log("\ncontextMissing_MultiOnCompletion");
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // run some database queries
                .collect().asList()
                // commit transaction if there was no error (onCompletion)
                .toMulti()
                .onCompletion().call(this::commitTransaction)
                .toUni()
                // some other async step
                .call((items) -> Uni.createFrom().item(items))
                .log();
    }

    public Uni<List<String>> contextMissing_MultiOnFailure() {
        log("\ncontextMissing_MultiOnFailure");
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // cause failure
                .invoke(item -> {
                    if (item.equals("b")) {
                        throw new RuntimeException("fail");
                    }
                })
                // run some database queries
                .collect().asList()
                // commit transaction if there was an error (onFailure)
                .toMulti()
                .onFailure()
                .call(this::commitTransaction)
                .toUni()
                // some other async step
                .call(items -> Uni.createFrom().item(items))
                .log();
    }

    public Uni<List<String>> contextAvailable_UniOnItem() {
        log("\ncontextAvailable_UniOnItem");
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // run some database queries
                .collect().asList()
                // commit transaction if there was no error
                .onItem().call(this::commitTransaction)
                // some other async step
                .call(items -> Uni.createFrom().item(items))
                .log();
    }

    public Uni<List<String>> contextMissing_UniOnTermination() {
        log("\ncontextMissing_UniOnCompletion");
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // run some database queries
                .collect().asList()
                // commit transaction in any case
                .onTermination().call(this::commitTransaction)
                // some other async step
                .call(items -> Uni.createFrom().item(items))
                .log();
    }

    public Uni<List<String>> contextAvailable_UniOnFailure() {
        log("\ncontextMissing_UniOnFailure");
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // cause failure
                .invoke(item -> {
                    if (item.equals("b")) {
                        throw new RuntimeException("fail");
                    }
                })
                // run some database queries
                .collect().asList()
                // rollback if there was an error
                .onFailure().call(this::rollbackTransaction)
                // some other async step
                .call(items -> Uni.createFrom().item(items))
                .log();
    }

    public Uni<Void> beginTransaction() {
        return Uni.createFrom().context(ctx -> {
            log("context when beginning:" + ctx);
            // open transaction (placeholder example)
            ctx.put("transaction", new Object());
            return Uni.createFrom().voidItem();
        });
    }

    public Uni<Void> commitTransaction() {
        return Uni.createFrom().context(ctx -> {
            log("context when committing:" + ctx);
            // commit transaction
            // example code:
            return Uni.createFrom().item(() -> ctx.get("transaction")).replaceWithVoid();
            // real code:
            //return Uni.createFrom().publisher(ctx.get("transaction").commit())
        });
    }

    public Uni<Void> rollbackTransaction() {
        return Uni.createFrom().context(ctx -> {
            log("context when rolling back:" + ctx);
            // rollback transaction
            // example code:
            return Uni.createFrom().item(() -> ctx.get("transaction")).replaceWithVoid();
            // real code:
            //return Uni.createFrom().publisher(ctx.get("transaction").rollback())
        });
    }

This is the output:

contextAvailable_MultiOnItem
[--> Uni.UniOnItemTransformToUni.0 | onSubscribe()
context when beginning:Context{entries=null}
Item: a
Item: b
Item: c
Item: d
context when committing:Context{entries={transaction=java.lang.Object@1a4b9d39}}
[--> Uni.UniOnItemTransformToUni.0 | onItem([a, b, c, d])

contextMissing_MultiOnCompletion
[--> Uni.UniOnItemTransformToUni.0 | onSubscribe()
context when beginning:Context{entries=null}
Item: a
Item: b
Item: c
Item: d
[--> Uni.UniOnItemTransformToUni.0 | onItem([a, b, c, d])
context when committing:Context{entries=null}

contextMissing_MultiOnFailure
[--> Uni.UniOnItemTransformToUni.0 | onSubscribe()
context when beginning:Context{entries=null}
Item: a
Item: b
context when committing:Context{entries=null}
[--> Uni.UniOnItemTransformToUni.0 | onFailure(io.smallrye.mutiny.CompositeException("Multiple exceptions caught:
	[Exception 0] java.lang.RuntimeException: fail
	[Exception 1] java.util.NoSuchElementException: The context is empty"))

contextAvailable_UniOnItem
[--> Uni.UniOnItemTransformToUni.0 | onSubscribe()
context when beginning:Context{entries=null}
Item: a
Item: b
Item: c
Item: d
context when committing:Context{entries={transaction=java.lang.Object@2b779686}}
[--> Uni.UniOnItemTransformToUni.0 | onItem([a, b, c, d])
[-- Mutiny had to drop the following exception --]
Exception received by: io.smallrye.mutiny.infrastructure.Infrastructure.handleDroppedException(Infrastructure.java:300)
java.util.NoSuchElementException: The context is empty
[ VERY LONG USELESS STACK TRACE OMITTED ]
[------------------------------------------------]

contextMissing_UniOnCompletion
[--> Uni.UniOnItemTransformToUni.0 | onSubscribe()
context when beginning:Context{entries=null}
Item: a
Item: b
Item: c
Item: d
context when committing:Context{entries=null}
[--> Uni.UniOnItemTransformToUni.0 | onFailure(java.util.NoSuchElementException("The context is empty"))

contextMissing_UniOnFailure
[--> Uni.UniOnItemTransformToUni.0 | onSubscribe()
context when beginning:Context{entries=null}
Item: a
Item: b
context when rolling back:Context{entries={transaction=java.lang.Object@42ea06c9}}
[--> Uni.UniOnItemTransformToUni.0 | onFailure(java.lang.RuntimeException("fail"))

I have had success with the following workaround. However, this requires "hiding" the choice of onTermination vs onFailure inside the commitTransactionCtx helper function. I'd like to be able to explicitly choose in the top-level pipeline.

    public <T> Multi<T> commitTransactionCtx(final Multi<T> multi, final Context ctx) {
        return multi.onTermination().call(() -> {
            log("commitTransactionCtx: " + ctx);
            // commit transaction ...
            return Uni.createFrom().voidItem();
        });
    }
    
        public Uni<List<String>> contextAvailable_MultiWithContextOnTermination() {
        return beginTransaction()
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                .invoke(item -> log("Item: " + item))
                // run some database queries
                .collect().asList()
                // commit transaction if there was no error (onTermination)
                .toMulti()
                .withContext(this::commitTransactionCtx)
                .toUni()
                // some other async step
                .call((items) -> Uni.createFrom().item(items))
                .log();
    }

[jponge]

The onFailure / onCompletion / onTermination groups are not much different to the other operators (see MultiOperatorProcessor which is the base class for many operators, and brings a default context passing).

I think you should try to come up with a unit test even smaller than that to possibly illustrate a faulty context passing, if any. The more atomic, the better.

[ivu-mawi]

Here are two unit tests: the one with onItem succeeds, the one with onCompletion fails. I would expect both to succeed.

package com.example.starter;

import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.concurrent.atomic.AtomicBoolean;

import io.smallrye.mutiny.Multi;
import io.smallrye.mutiny.Uni;
import io.smallrye.mutiny.helpers.test.AssertSubscriber;
import io.smallrye.mutiny.helpers.test.UniAssertSubscriber;

import org.junit.jupiter.api.Test;


class ContextPassingTest {
    public Uni<Void> beginTransaction(final AtomicBoolean isTransactionCommitted) {
        return Uni.createFrom().context(ctx -> {
            // open transaction (placeholder boolean)
            System.out.println("Opening transaction");
            ctx.put("transaction", isTransactionCommitted);
            return Uni.createFrom().voidItem();
        });
    }

    public Uni<Void> commitTransaction() {
        return Uni.createFrom().context(ctx -> {
            return Uni.createFrom().item(() -> {
                System.out.println("Going to commit transaction");
                if (ctx.contains("transaction")) {
                    // commit transaction (placeholder boolean)
                    ctx.<AtomicBoolean>get("transaction").set(true);
                }
                return null;
            }).replaceWithVoid();
        });
    }

    @Test
    void contextPassing_multiOnCompletion() {
        final var isTransactionCommitted = new AtomicBoolean(false);
        final var multi = beginTransaction(isTransactionCommitted)
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                // do some other db queries
                // ...
                // commit transaction if there was no error (onCompletion)
                .onCompletion().call(this::commitTransaction);

        multi.subscribe().withSubscriber(AssertSubscriber.create(10)).assertCompleted();
        assertTrue(isTransactionCommitted.get());
    }

    @Test
    void contextPassing_uniOnItem() {
        final var transactionCommitted = new AtomicBoolean(false);
        final var uni = beginTransaction(transactionCommitted)
                // get items from db
                .onItem().transformToMulti(ignore -> Multi.createFrom().items("a", "b", "c", "d"))
                // run some database queries
                // ...
                .collect().asList()
                // commit transaction if there was no error (onItem)
                .onItem().call(this::commitTransaction);

        uni.subscribe().withSubscriber(UniAssertSubscriber.create()).assertCompleted();
        assertTrue(transactionCommitted.get());
    }
}