POC for context capturing in loom

CapturedScopedValueContext.java

import java.util.concurrent.Callable;

public interface CapturedScopedValueContext {
    <T> T inContext(Callable<? extends T> task);
}

CapturedScopes.java

import java.util.Objects;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.function.Consumer;
import jdk.incubator.concurrent.StructuredTaskScope;
import org.jtrim2.cancel.Cancellation;
import org.jtrim2.collections.ReservablePollingQueues;
import org.jtrim2.concurrent.collections.TerminableQueues;
import org.jtrim2.concurrent.collections.TerminatedQueueException;

public final class CapturedScopes {
    public static void withCurrentContext(
            Consumer<? super CapturedScopedValueContext> task
    ) {
        Objects.requireNonNull(task, "task");

        try (var scope = new StructuredTaskScope.ShutdownOnFailure()) {
            var queue = TerminableQueues.<Runnable>withWrappedQueue(
                    ReservablePollingQueues.createFifoQueue(1)
            );
            try {
                scope.fork(() -> {
                    try {
                        while (true) {
                            var receivedTask = queue.take(Cancellation.UNCANCELABLE_TOKEN);
                            scope.fork(() -> {
                                receivedTask.run();
                                return null;
                            });
                        }
                    } catch (TerminatedQueueException e) {
                        // Fine
                    }
                    return null;
                });
                task.accept(new CapturedScopedValueContext() {
                    @Override
                    public <T> T inContext(Callable<? extends T> task) {
                        try {
                            var result = new CompletableFuture<T>();
                            queue.put(Cancellation.UNCANCELABLE_TOKEN, () -> {
                                try {
                                    result.complete(task.call());
                                } catch (Throwable ex) {
                                    result.completeExceptionally(ex);
                                }
                            });
                            return result.get();
                        } catch (InterruptedException
                                 | ExecutionException
                                 | TerminatedQueueException e
                        ) {
                            throw new RuntimeException(e);
                        }
                    }
                });
            } finally {
                queue.shutdownAndWaitUntilEmpty(Cancellation.UNCANCELABLE_TOKEN);
            }
            scope.join().throwIfFailed();
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    private CapturedScopes() {
        throw new AssertionError();
    }
}

ExampleInjector.java

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.function.Supplier;
import jdk.incubator.concurrent.ScopedValue;

public final class ExampleInjector {
    private static final ScopedValue<Map<Class<?>, Supplier<?>>> FACTORIES
            = ScopedValue.newInstance();

    private static <T> void withBindingUnsafe(
            Class<T> type,
            Supplier<T> supplier,
            Runnable task
    ) {
        var prevValue = FACTORIES.orElse(Collections.emptyMap());
        var newValue = new HashMap<>(prevValue);
        newValue.put(type, supplier);

        ScopedValue.where(FACTORIES, newValue, task);
    }

    public static <T> T getByType(Class<T> type) {
        Objects.requireNonNull(type, "type");

        Supplier<?> provider = FACTORIES
                .orElse(Collections.emptyMap())
                .get(type);
        if (provider == null) {
            throw new IllegalStateException();
        }
        return type.cast(provider.get());
    }

    public static <T> void withBinding(
            Class<T> type,
            Supplier<T> supplier,
            Runnable task
    ) {
        Objects.requireNonNull(type, "type");
        Objects.requireNonNull(supplier, "supplier");
        Objects.requireNonNull(task, "task");

        CapturedScopes.withCurrentContext(context -> {
            withBindingUnsafe(type, () -> context.inContext(supplier::get), task);
        });
    }

    private ExampleInjector() {
        throw new AssertionError();
    }
}

InjectorTest.java

import jdk.incubator.concurrent.ScopedValue;

public final class InjectorTest {
    public static void main(String[] args) {
        ScopedValue<String> testValue = ScopedValue.newInstance();

        ScopedValue.where(testValue, "OuterValue", () -> {
            ExampleInjector.withBinding(String.class, testValue::get, () -> {
                ScopedValue.where(testValue, "InnerValue", () -> {
                    System.out.println(ExampleInjector.getByType(String.class));
                });
            });
        });
    }
}