Cleanup

Author: bryce-anderson

servicetalk-benchmarks/src/jmh/java/io/servicetalk/concurrent/api/AsyncContextProviderBenchmark.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright © 2025 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.servicetalk.concurrent.api;
+
+import io.servicetalk.context.api.ContextMap;
+
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+
+/**
+ *
+ */
+@Fork(1)
+@State(Scope.Benchmark)
+@Warmup(iterations = 5, time = 3)
+@Measurement(iterations = 5, time = 3)
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@BenchmarkMode(Mode.AverageTime)
+public class AsyncContextProviderBenchmark {
+
+    /**
+     * gc profiling of the DefaultAsyncContextProvider shows that the Scope based detachment can be stack allocated
+     * at least under some conditions.
+     *
+     * Benchmark                                                                   Mode  Cnt   Score    Error   Units
+     * AsyncContextProviderBenchmark.contextRestoreCost                            avgt    5   3.932 ±  0.022   ns/op
+     * AsyncContextProviderBenchmark.contextRestoreCost:gc.alloc.rate              avgt    5  ≈ 10⁻⁴           MB/sec
+     * AsyncContextProviderBenchmark.contextRestoreCost:gc.alloc.rate.norm         avgt    5  ≈ 10⁻⁶             B/op
+     * AsyncContextProviderBenchmark.contextRestoreCost:gc.count                   avgt    5     ≈ 0           counts
+     * AsyncContextProviderBenchmark.contextSaveAndRestoreCost                     avgt    5   1.712 ±  0.005   ns/op
+     * AsyncContextProviderBenchmark.contextSaveAndRestoreCost:gc.alloc.rate       avgt    5  ≈ 10⁻⁴           MB/sec
+     * AsyncContextProviderBenchmark.contextSaveAndRestoreCost:gc.alloc.rate.norm  avgt    5  ≈ 10⁻⁷             B/op
+     * AsyncContextProviderBenchmark.contextSaveAndRestoreCost:gc.count            avgt    5     ≈ 0           counts
+     */
+
+    private static final ContextMap.Key<String> KEY = ContextMap.Key.newKey("test-key", String.class);
+    private static final String EXPECTED = "hello, world!";
+
+    private static Function<String, String> wrappedFunction;
+
+    @Setup
+    public void setup() {
+        // This will capture the current context
+        wrappedFunction = AsyncContext.wrapFunction(ignored -> AsyncContext.context().get(KEY));
+        AsyncContext.context().put(KEY, EXPECTED);
+    }
+
+    @Benchmark
+    public String contextRestoreCost() {
+        return wrappedFunction.apply("ignored");
+    }
+
+    @Benchmark
+    public String contextSaveAndRestoreCost() {
+        return AsyncContext.wrapFunction(Function.<String>identity()).apply("ignored");
+    }
+}

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

@@ -72,11 +72,10 @@ public final class AsyncContext {
 
     static {
         AsyncContextProvider result = DefaultAsyncContextProvider.INSTANCE;
-        for (UnaryOperator<AsyncContextProvider> wrapper : asyncProviderWrappers()) {
-            System.out.println("Wrapping with " + wrapper.getClass().getName());
+        List<UnaryOperator<AsyncContextProvider>> wrappers = asyncProviderWrappers();
+        for (UnaryOperator<AsyncContextProvider> wrapper : wrappers) {
             result = wrapper.apply(result);
         }
-        System.out.println("Default AsyncContextProvider: " + result.getClass().getName());
         DEFAULT_ENABLED_PROVIDER = result;
         provider = DEFAULT_ENABLED_PROVIDER;
     }

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

@@ -36,7 +36,7 @@
 /**
  * Implementation that backs the {@link AsyncContext}.
  */
-public interface AsyncContextProvider {
+interface AsyncContextProvider {
     /**
      * Get the current context.
      *

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

@@ -1730,7 +1730,7 @@ public final Future<Void> toFuture() {
      */
     ContextMap contextForSubscribe(AsyncContextProvider provider) {
         // the default behavior is to copy the map. Some operators may want to use shared map
-        return provider.saveContext().copy(); // TODO: should we have a `saveContextCopy()` for perf reasons?
+        return provider.saveContext().copy();
     }
 
     /**
@@ -2266,12 +2266,14 @@ private void subscribeWithContext(Subscriber subscriber,
                                       AsyncContextProvider contextProvider, ContextMap contextMap) {
         requireNonNull(subscriber);
         Subscriber wrapped = contextProvider.wrapCancellable(subscriber, contextMap);
-        if (contextProvider.context() == contextMap) { // TODO: this shortcut will not work anymore if `saveContext()` returns a result different than `context()`. Maybe a provider.isActive(context) would work?
+        if (contextProvider.context() == contextMap) {
             // No need to wrap as we are sharing the AsyncContext
             handleSubscribe(wrapped, contextMap, contextProvider);
         } else {
             // Ensure that AsyncContext used for handleSubscribe() is the contextMap for the subscribe()
-            contextProvider.wrapRunnable(() -> handleSubscribe(wrapped, contextMap, contextProvider), contextMap).run();
+            try(Scope unused = contextProvider.attachContext(contextMap)) {
+                handleSubscribe(wrapped, contextMap, contextProvider);
+            }
         }
     }
 

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

@@ -4867,7 +4867,9 @@ private void subscribeWithContext(Subscriber<? super T> subscriber,
             handleSubscribe(wrapped, contextMap, provider);
         } else {
             // Ensure that AsyncContext used for handleSubscribe() is the contextMap for the subscribe()
-            provider.wrapRunnable(() -> handleSubscribe(wrapped, contextMap, provider), contextMap).run();
+            try (Scope ignored = provider.attachContext(contextMap)) {
+                handleSubscribe(wrapped, contextMap, provider);
+            }
         }
     }
 

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

@@ -1,7 +1,35 @@
+/*
+ * Copyright © 2025 Apple Inc. and the ServiceTalk project authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
 package io.servicetalk.concurrent.api;
 
+import io.servicetalk.context.api.ContextMap;
+
+/**
+ * An abstraction for detaching a context from the current thread.
+ *
+ * This abstraction is intended to allow the modifications performed by
+ * {@link AsyncContextProvider#attachContext(ContextMap)} to be undone. In practice, this may look like restoring
+ * a {@link ThreadLocal} to the state it had before the call to {@link AsyncContextProvider#attachContext(ContextMap)}
+ * call.
+ */
 public interface Scope extends AutoCloseable {
 
+    /**
+     * No-op {@link Scope}.
+     */
     Scope NOOP = () -> {};
 
     @Override

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

@@ -2711,7 +2711,9 @@ private void subscribeWithContext(Subscriber<? super T> subscriber,
             handleSubscribe(wrapped, contextMap, contextProvider);
         } else {
             // Ensure that AsyncContext used for handleSubscribe() is the contextMap for the subscribe()
-            contextProvider.wrapRunnable(() -> handleSubscribe(wrapped, contextMap, contextProvider), contextMap).run();
+            try (Scope ignored = contextProvider.attachContext(contextMap)) {
+                handleSubscribe(wrapped, contextMap, contextProvider);
+            }
         }
     }
 

servicetalk-opentelemetry-http/src/main/java/io/servicetalk/opentelemetry/http/OpenTelemetryHttpRequestFilter.java

@@ -172,7 +172,6 @@ public Single<StreamingHttpResponse> request(final StreamingHttpRequest request)
     private Single<StreamingHttpResponse> trackRequest(final StreamingHttpRequester delegate,
                                                        final StreamingHttpRequest request) {
         final Context parentContext = Context.current();
-        System.out.println(OpenTelemetryHttpRequestFilter.class.getSimpleName() + ": Current context: " + parentContext + ", thread: " + Thread.currentThread());
         final Context context = instrumenter.start(parentContext, request);
 
         final Scope scope = context.makeCurrent();