concurrent-api: make context capture more generic

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractAsynchronousCompletableOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -38,13 +36,13 @@ abstract class AbstractAsynchronousCompletableOperator extends AbstractNoHandleS
 
     @Override
     final void handleSubscribe(Subscriber subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
         // The AsyncContext needs to be preserved when ever we interact with the original Subscriber, so we wrap it here
         // with the original contextMap. Otherwise some other context may leak into this subscriber chain from the other
         // side of the asynchronous boundary.
         final Subscriber operatorSubscriber =
-                contextProvider.wrapCompletableSubscriberAndCancellable(subscriber, contextMap);
+                contextProvider.wrapCompletableSubscriberAndCancellable(subscriber, capturedContext);
         final Subscriber upstreamSubscriber = apply(operatorSubscriber);
-        original.delegateSubscribe(upstreamSubscriber, contextMap, contextProvider);
+        original.delegateSubscribe(upstreamSubscriber, capturedContext, contextProvider);
     }
 }

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractAsynchronousPublisherOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -41,13 +39,13 @@ abstract class AbstractAsynchronousPublisherOperator<T, R> extends AbstractNoHan
 
     @Override
     final void handleSubscribe(Subscriber<? super R> subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
         // The AsyncContext needs to be preserved when ever we interact with the original Subscriber, so we wrap it here
         // with the original contextMap. Otherwise some other context may leak into this subscriber chain from the other
         // side of the asynchronous boundary.
         final Subscriber<? super R> operatorSubscriber =
-                contextProvider.wrapPublisherSubscriberAndSubscription(subscriber, contextMap);
+                contextProvider.wrapPublisherSubscriberAndSubscription(subscriber, capturedContext);
         final Subscriber<? super T> upstreamSubscriber = apply(operatorSubscriber);
-        original.delegateSubscribe(upstreamSubscriber, contextMap, contextProvider);
+        original.delegateSubscribe(upstreamSubscriber, capturedContext, contextProvider);
     }
 }

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractAsynchronousSingleOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -41,13 +39,13 @@ abstract class AbstractAsynchronousSingleOperator<T, R> extends AbstractNoHandle
 
     @Override
     final void handleSubscribe(Subscriber<? super R> subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
         // The AsyncContext needs to be preserved when ever we interact with the original Subscriber, so we wrap it here
         // with the original contextMap. Otherwise some other context may leak into this subscriber chain from the other
         // side of the asynchronous boundary.
         final Subscriber<? super R> operatorSubscriber =
-                contextProvider.wrapSingleSubscriberAndCancellable(subscriber, contextMap);
+                contextProvider.wrapSingleSubscriberAndCancellable(subscriber, capturedContext);
         final Subscriber<? super T> upstreamSubscriber = apply(operatorSubscriber);
-        original.delegateSubscribe(upstreamSubscriber, contextMap, contextProvider);
+        original.delegateSubscribe(upstreamSubscriber, capturedContext, contextProvider);
     }
 }

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractCompletableAndSingleConcatenated.java

@@ -18,21 +18,21 @@
 import io.servicetalk.concurrent.Cancellable;
 import io.servicetalk.concurrent.CompletableSource;
 import io.servicetalk.concurrent.internal.SequentialCancellable;
-import io.servicetalk.context.api.ContextMap;
 
 import javax.annotation.Nullable;
 
 abstract class AbstractCompletableAndSingleConcatenated<T> extends AbstractNoHandleSubscribeSingle<T> {
 
     @Override
     protected void handleSubscribe(final Subscriber<? super T> subscriber,
-                                   final ContextMap contextMap, final AsyncContextProvider contextProvider) {
-        final Subscriber<? super T> wrappedSubscriber = contextProvider.wrapSingleSubscriber(subscriber, contextMap);
-        delegateSubscribeToOriginal(wrappedSubscriber, contextMap, contextProvider);
+                                   final CapturedContext capturedContext, final AsyncContextProvider contextProvider) {
+        final Subscriber<? super T> wrappedSubscriber =
+                contextProvider.wrapSingleSubscriber(subscriber, capturedContext);
+        delegateSubscribeToOriginal(wrappedSubscriber, capturedContext, contextProvider);
     }
 
     abstract void delegateSubscribeToOriginal(Subscriber<? super T> offloadSubscriber,
-                                              ContextMap contextMap, AsyncContextProvider contextProvider);
+                                              CapturedContext capturedContext, AsyncContextProvider contextProvider);
 
     abstract static class AbstractConcatWithSubscriber<T> implements Subscriber<T>, CompletableSource.Subscriber {
 

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractMergeCompletableOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 abstract class AbstractMergeCompletableOperator<T extends CompletableMergeSubscriber>
@@ -30,15 +28,15 @@ abstract class AbstractMergeCompletableOperator<T extends CompletableMergeSubscr
     }
 
     @Override
-    final void handleSubscribe(Subscriber subscriber, ContextMap contextMap,
+    final void handleSubscribe(Subscriber subscriber, CapturedContext capturedContext,
                                AsyncContextProvider contextProvider) {
         // The AsyncContext needs to be preserved when ever we interact with the original Subscriber, so we wrap it here
         // with the original contextMap. Otherwise some other context may leak into this subscriber chain from the other
         // side of the asynchronous boundary.
         final Subscriber operatorSubscriber =
-                contextProvider.wrapCompletableSubscriberAndCancellable(subscriber, contextMap);
+                contextProvider.wrapCompletableSubscriberAndCancellable(subscriber, capturedContext);
         T mergeSubscriber = apply(operatorSubscriber);
-        original.delegateSubscribe(mergeSubscriber, contextMap, contextProvider);
+        original.delegateSubscribe(mergeSubscriber, capturedContext, contextProvider);
         doMerge(mergeSubscriber);
     }
 

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractPubToCompletable.java

@@ -17,7 +17,6 @@
 
 import io.servicetalk.concurrent.PublisherSource;
 import io.servicetalk.concurrent.internal.DelayedCancellable;
-import io.servicetalk.context.api.ContextMap;
 
 /**
  * A {@link Completable} created from a {@link Publisher}.
@@ -40,12 +39,12 @@ abstract class AbstractPubToCompletable<T> extends AbstractNoHandleSubscribeComp
 
     @Override
     final void handleSubscribe(final Subscriber subscriber,
-                               final ContextMap contextMap, final AsyncContextProvider contextProvider) {
+                               final CapturedContext capturedContext, final AsyncContextProvider contextProvider) {
         // We are now subscribing to the original Publisher chain for the first time, wrap Subscription to preserve the
         // context.
         PublisherSource.Subscriber<? super T> wrappedSubscriber =
-                contextProvider.wrapSubscription(newSubscriber(subscriber), contextMap);
-        source.delegateSubscribe(wrappedSubscriber, contextMap, contextProvider);
+                contextProvider.wrapSubscription(newSubscriber(subscriber), capturedContext);
+        source.delegateSubscribe(wrappedSubscriber, capturedContext, contextProvider);
     }
 
     abstract static class AbstractPubToCompletableSubscriber<T> extends DelayedCancellable

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractPubToSingle.java

@@ -17,7 +17,6 @@
 
 import io.servicetalk.concurrent.PublisherSource;
 import io.servicetalk.concurrent.PublisherSource.Subscription;
-import io.servicetalk.context.api.ContextMap;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -38,12 +37,12 @@ abstract class AbstractPubToSingle<T> extends AbstractNoHandleSubscribeSingle<T>
 
     @Override
     final void handleSubscribe(final Subscriber<? super T> subscriber,
-                               final ContextMap contextMap, final AsyncContextProvider contextProvider) {
+                               final CapturedContext capturedContext, final AsyncContextProvider contextProvider) {
         // We are now subscribing to the original Publisher chain for the first time, wrap Subscription to preserve the
         // context.
         PublisherSource.Subscriber<? super T> wrappedSubscription =
-                contextProvider.wrapSubscription(newSubscriber(subscriber), contextMap);
-        source.delegateSubscribe(wrappedSubscription, contextMap, contextProvider);
+                contextProvider.wrapSubscription(newSubscriber(subscriber), capturedContext);
+        source.delegateSubscribe(wrappedSubscription, capturedContext, contextProvider);
     }
 
     abstract PublisherSource.Subscriber<T> newSubscriber(Subscriber<? super T> original);

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractPublisherGroupBy.java

@@ -18,7 +18,6 @@
 import io.servicetalk.concurrent.PublisherSource;
 import io.servicetalk.concurrent.internal.ConcurrentSubscription;
 import io.servicetalk.concurrent.internal.DuplicateSubscribeException;
-import io.servicetalk.context.api.ContextMap;
 
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
@@ -49,21 +48,21 @@ abstract class AbstractPublisherGroupBy<Key, T> extends AbstractNoHandleSubscrib
     abstract static class AbstractGroupBySubscriber<Key, T> implements Subscriber<T> {
         private boolean rootCancelled;
         private final int queueLimit;
-        private final ContextMap contextMap;
+        private final CapturedContext capturedContext;
         private final AsyncContextProvider contextProvider;
         private final Map<Key, GroupMulticastSubscriber<Key, T>> groups;
         private final GroupMulticastSubscriber<String, GroupedPublisher<Key, T>> target;
         @Nullable
         private Subscription subscription;
 
         AbstractGroupBySubscriber(final Subscriber<? super GroupedPublisher<Key, T>> target, final int queueLimit,
-                                  final int initialCapacityForGroups, final ContextMap contextMap,
+                                  final int initialCapacityForGroups, final CapturedContext capturedContext,
                                   final AsyncContextProvider contextProvider) {
             this.queueLimit = queueLimit;
-            this.contextMap = contextMap;
+            this.capturedContext = capturedContext;
             this.contextProvider = contextProvider;
             this.target = new GroupMulticastSubscriber<>(this, "root");
-            this.target.subscriber(target, false, contextMap, contextProvider);
+            this.target.subscriber(target, false, capturedContext, contextProvider);
             groups = new ConcurrentHashMap<>(initialCapacityForGroups);
         }
 
@@ -98,7 +97,7 @@ final void onNext(Key key, @Nullable T t) {
             } else {
                 groupSub = new GroupMulticastSubscriber<>(this, key);
                 GroupedPublisher<Key, T> groupedPublisher = new DefaultGroupedPublisher<>(key, groupSub,
-                        contextMap, contextProvider);
+                        capturedContext, contextProvider);
                 final GroupMulticastSubscriber<Key, T> oldVal = groups.put(key, groupSub);
                 assert oldVal == null; // concurrent onNext not allowed, collision not expected.
                 groupSub.onNext(t); // deliver to group first to avoid re-entry creating ordering issues.
@@ -165,16 +164,16 @@ public String toString() {
         }
 
         void subscriber(final Subscriber<? super T> subscriber, final boolean triggerOnSubscribe,
-                        final ContextMap contextMap, final AsyncContextProvider contextProvider) {
+                        final CapturedContext capturedContext, final AsyncContextProvider contextProvider) {
             // The root Subscriber's downstream subscriber is set internally, so no need for atomic operation to filter
             // duplicates.
             if (!triggerOnSubscribe) {
                 assert this.subscriber == null && ctxSubscriber == null;
                 this.subscriber = subscriber;
-                ctxSubscriber = contextProvider.wrapPublisherSubscriber(subscriber, contextMap);
+                ctxSubscriber = contextProvider.wrapPublisherSubscriber(subscriber, capturedContext);
             } else if (subscriberStateUpdater.compareAndSet(this, 0, 1)) {
                 this.subscriber = subscriber;
-                ctxSubscriber = contextProvider.wrapPublisherSubscriber(subscriber, contextMap);
+                ctxSubscriber = contextProvider.wrapPublisherSubscriber(subscriber, capturedContext);
                 triggerOnSubscribe();
             } else {
                 // this.subscriber may be null (we set the subscriber variable after subscriberStateUpdater),
@@ -215,14 +214,14 @@ int outstandingDemandLimit() {
     private static final class DefaultGroupedPublisher<Key, T> extends GroupedPublisher<Key, T>
             implements PublisherSource<T> {
         private final GroupMulticastSubscriber<Key, T> groupSink;
-        private final ContextMap contextMap;
+        private final CapturedContext capturedContext;
         private final AsyncContextProvider contextProvider;
 
         DefaultGroupedPublisher(final Key key, final GroupMulticastSubscriber<Key, T> groupSink,
-                                final ContextMap contextMap, final AsyncContextProvider contextProvider) {
+                                final CapturedContext capturedContext, final AsyncContextProvider contextProvider) {
             super(key);
             this.groupSink = groupSink;
-            this.contextMap = contextMap;
+            this.capturedContext = capturedContext;
             this.contextProvider = contextProvider;
         }
 
@@ -233,7 +232,7 @@ public void subscribe(final Subscriber<? super T> subscriber) {
 
         @Override
         protected void handleSubscribe(Subscriber<? super T> sub) {
-            groupSink.subscriber(sub, true, contextMap, contextProvider);
+            groupSink.subscriber(sub, true, capturedContext, contextProvider);
         }
     }
 }

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractSynchronousCompletable.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 /**
  * Base class for all {@link Completable}s that are created with already realized result and does not generate result
  * asynchronously.
@@ -25,10 +23,10 @@ abstract class AbstractSynchronousCompletable extends AbstractNoHandleSubscribeC
 
     @Override
     final void handleSubscribe(Subscriber subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
         // We need to wrap the Subscriber to save/restore the AsyncContext on each operation or else the AsyncContext
         // may leak from another thread.
-        doSubscribe(contextProvider.wrapCompletableSubscriber(subscriber, contextMap));
+        doSubscribe(contextProvider.wrapCompletableSubscriber(subscriber, capturedContext));
     }
 
     /**

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractSynchronousCompletableOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -41,7 +39,7 @@ abstract class AbstractSynchronousCompletableOperator extends AbstractNoHandleSu
 
     @Override
     final void handleSubscribe(Subscriber subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
-        original.delegateSubscribe(apply(subscriber), contextMap, contextProvider);
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
+        original.delegateSubscribe(apply(subscriber), capturedContext, contextProvider);
     }
 }

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractSynchronousPublisher.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 /**
  * Base class for all {@link Publisher}s that are created with already realized values and do not generate values
  * asynchronously.
@@ -27,10 +25,10 @@ abstract class AbstractSynchronousPublisher<T> extends AbstractNoHandleSubscribe
 
     @Override
     final void handleSubscribe(Subscriber<? super T> subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
         // We need to wrap the Subscriber to save/restore the AsyncContext on each operation or else the AsyncContext
         // may leak from another thread.
-        doSubscribe(contextProvider.wrapPublisherSubscriber(subscriber, contextMap));
+        doSubscribe(contextProvider.wrapPublisherSubscriber(subscriber, capturedContext));
     }
 
     /**

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractSynchronousPublisherOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -43,7 +41,7 @@ abstract class AbstractSynchronousPublisherOperator<T, R> extends AbstractNoHand
 
     @Override
     final void handleSubscribe(Subscriber<? super R> subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
-        original.delegateSubscribe(apply(subscriber), contextMap, contextProvider);
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
+        original.delegateSubscribe(apply(subscriber), capturedContext, contextProvider);
     }
 }

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractSynchronousSingle.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 /**
  * Base class for all {@link Single}s that are created with already realized result and does not generate result
  * asynchronously.
@@ -27,10 +25,10 @@ abstract class AbstractSynchronousSingle<T> extends AbstractNoHandleSubscribeSin
 
     @Override
     final void handleSubscribe(Subscriber<? super T> subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
         // We need to wrap the Subscriber to save/restore the AsyncContext on each operation or else the AsyncContext
         // may leak from another thread.
-        doSubscribe(contextProvider.wrapSingleSubscriber(subscriber, contextMap));
+        doSubscribe(contextProvider.wrapSingleSubscriber(subscriber, capturedContext));
     }
 
     /**

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/AbstractSynchronousSingleOperator.java

@@ -15,8 +15,6 @@
  */
 package io.servicetalk.concurrent.api;
 
-import io.servicetalk.context.api.ContextMap;
-
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -43,7 +41,7 @@ abstract class AbstractSynchronousSingleOperator<T, R> extends AbstractNoHandleS
 
     @Override
     final void handleSubscribe(Subscriber<? super R> subscriber,
-                               ContextMap contextMap, AsyncContextProvider contextProvider) {
-        original.delegateSubscribe(apply(subscriber), contextMap, contextProvider);
+                               CapturedContext capturedContext, AsyncContextProvider contextProvider) {
+        original.delegateSubscribe(apply(subscriber), capturedContext, contextProvider);
     }
 }