concurrent-api: shave some allocations from new Scope (#3185)

Motivation:

In the allocation profile we see anonymous instances of the new `Scope` being
created. While not a big deal, we can do better.

Modifications:

We can shave these off of the default implementation by making `CopyOnWriteContextMap`,
and really any context map, an instance of `Scope`. The `close()` method just needs
to then ask the current provider to set the context to itself.

Result:

Less allocations.

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

@@ -21,6 +21,7 @@
 import io.servicetalk.concurrent.PublisherSource.Subscription;
 import io.servicetalk.concurrent.SingleSource;
 import io.servicetalk.context.api.ContextMap;
+import io.servicetalk.context.api.ContextMapHolder;
 
 import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
@@ -36,7 +37,7 @@
 /**
  * Implementation that backs the {@link AsyncContext}.
  */
-interface AsyncContextProvider {
+interface AsyncContextProvider extends ContextMapHolder {
     /**
      * Get the current context.
      *
@@ -45,6 +46,7 @@ interface AsyncContextProvider {
      *
      * @return The current context.
      */
+    @Override
     ContextMap context();
 
     /**

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

@@ -39,7 +39,7 @@
  * {@link ContextMap.Key}-value entries in a single {@link ContextMap}. Common {@link ContextMap.Key}-value entries are
  * (tracing, MDC, auth, 3-custom user entries).
  */
-final class CopyOnWriteContextMap implements ContextMap {
+final class CopyOnWriteContextMap implements ContextMap, Scope {
     private static final AtomicReferenceFieldUpdater<CopyOnWriteContextMap, CopyContextMap> mapUpdater =
             AtomicReferenceFieldUpdater.newUpdater(CopyOnWriteContextMap.class, CopyContextMap.class, "map");
 
@@ -187,6 +187,11 @@ public String toString() {
         return ContextMapUtils.toString(this);
     }
 
+    @Override
+    public void close() {
+        AsyncContext.provider().context(this);
+    }
+
     private interface CopyContextMap {
 
         int size();

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

@@ -34,15 +34,18 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 import static java.lang.ThreadLocal.withInitial;
 
 final class DefaultAsyncContextProvider implements AsyncContextProvider {
 
-    private static final Logger LOGGER = LoggerFactory.getLogger(DefaultAsyncContextProvider.class);
     private static final ThreadLocal<ContextMap> CONTEXT_THREAD_LOCAL =
             withInitial(DefaultAsyncContextProvider::newContextMap);
 
+    private static final Logger LOGGER = LoggerFactory.getLogger(DefaultAsyncContextProvider.class);
+    private static final boolean NOT_IS_DEBUG_ENABLED = !LOGGER.isDebugEnabled();
+
     static final AsyncContextProvider INSTANCE = new DefaultAsyncContextProvider();
 
     private DefaultAsyncContextProvider() {
@@ -67,51 +70,28 @@ public ContextMap context() {
     }
 
     @Override
-    public ContextMap captureContext() {
-        return context();
-    }
-
-    @Override
-    public Scope attachContext(ContextMap contextMap) {
+    public ContextMapHolder context(@Nullable ContextMap contextMap) {
         final Thread currentThread = Thread.currentThread();
         if (currentThread instanceof ContextMapHolder) {
             final ContextMapHolder asyncContextMapHolder = (ContextMapHolder) currentThread;
-            ContextMap prev = asyncContextMapHolder.context();
             asyncContextMapHolder.context(contextMap);
-            return () -> detachContext(contextMap, prev == null ? newContextMap() : prev);
+        } else if (contextMap == null) {
+            CONTEXT_THREAD_LOCAL.remove();
         } else {
-            return slowPathSetContext(contextMap);
+            CONTEXT_THREAD_LOCAL.set(contextMap);
         }
+        return this;
     }
 
-    private static Scope slowPathSetContext(ContextMap contextMap) {
-        ContextMap prev = CONTEXT_THREAD_LOCAL.get();
-        CONTEXT_THREAD_LOCAL.set(contextMap);
-        return () -> detachContext(contextMap, prev);
-    }
-
-    private static void detachContext(ContextMap expectedContext, ContextMap toRestore) {
-        final Thread currentThread = Thread.currentThread();
-        if (currentThread instanceof ContextMapHolder) {
-            final ContextMapHolder asyncContextMapHolder = (ContextMapHolder) currentThread;
-            ContextMap current = asyncContextMapHolder.context();
-            if (current != expectedContext) {
-                LOGGER.debug("Current context didn't match the expected context. current: {}, expected: {}",
-                        current, expectedContext);
-            }
-            asyncContextMapHolder.context(toRestore);
-        } else {
-            slowPathDetachContext(expectedContext, toRestore);
-        }
+    @Override
+    public ContextMap captureContext() {
+        return context();
     }
 
-    private static void slowPathDetachContext(ContextMap expectedContext, ContextMap toRestore) {
-        ContextMap current = CONTEXT_THREAD_LOCAL.get();
-        if (current != expectedContext) {
-            LOGGER.debug("Current context didn't match the expected context. current: {}, expected: {}",
-                    current, expectedContext);
-        }
-        CONTEXT_THREAD_LOCAL.set(toRestore);
+    @Override
+    public Scope attachContext(ContextMap contextMap) {
+        ContextMap prev = exchangeContext(contextMap);
+        return NOT_IS_DEBUG_ENABLED && prev instanceof Scope ? (Scope) prev : () -> detachContext(contextMap, prev);
     }
 
     @Override
@@ -333,6 +313,31 @@ public <T, U, V> BiFunction<T, U, V> wrapBiFunction(final BiFunction<T, U, V> fu
         return new ContextPreservingBiFunction<>(func, context);
     }
 
+    private static ContextMap exchangeContext(ContextMap contextMap) {
+        ContextMap result;
+        final Thread currentThread = Thread.currentThread();
+        if (currentThread instanceof ContextMapHolder) {
+            final ContextMapHolder asyncContextMapHolder = (ContextMapHolder) currentThread;
+            result = asyncContextMapHolder.context();
+            if (result == null) {
+                result = newContextMap();
+            }
+            asyncContextMapHolder.context(contextMap);
+        } else {
+            result = CONTEXT_THREAD_LOCAL.get();
+            CONTEXT_THREAD_LOCAL.set(contextMap);
+        }
+        return result;
+    }
+
+    private static void detachContext(ContextMap expectedContext, ContextMap toRestore) {
+        ContextMap current = exchangeContext(toRestore);
+        if (current != expectedContext) {
+            LOGGER.debug("Current context didn't match the expected context. current: {}, expected: {}",
+                    current, expectedContext);
+        }
+    }
+
     private static ContextMap newContextMap() {
         return new CopyOnWriteContextMap();
     }

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

@@ -20,6 +20,7 @@
 import io.servicetalk.concurrent.SingleSource;
 import io.servicetalk.concurrent.internal.ContextMapUtils;
 import io.servicetalk.context.api.ContextMap;
+import io.servicetalk.context.api.ContextMapHolder;
 
 import java.util.Map;
 import java.util.concurrent.Callable;
@@ -56,6 +57,11 @@ public Scope attachContext(ContextMap contextMap) {
         return Scope.NOOP;
     }
 
+    @Override
+    public ContextMapHolder context(@Nullable ContextMap contextMap) {
+        return this;
+    }
+
     @Override
     public CompletableSource.Subscriber wrapCancellable(final CompletableSource.Subscriber subscriber,
                                                         final ContextMap context) {
@@ -165,7 +171,7 @@ public <T, U, V> BiFunction<T, U, V> wrapBiFunction(final BiFunction<T, U, V> fu
         return func;
     }
 
-    private static final class NoopContextMap implements ContextMap {
+    private static final class NoopContextMap implements ContextMap, Scope {
         static final ContextMap INSTANCE = new NoopContextMap();
 
         private NoopContextMap() {
@@ -280,5 +286,10 @@ public boolean equals(final Object o) {
         public String toString() {
             return ContextMapUtils.toString(this);
         }
+
+        @Override
+        public void close() {
+            // noop
+        }
     }
 }