Support micrometer context-propagation

core/build.gradle

@@ -108,6 +108,7 @@ dependencies {
     optionalApi libs.micrometer.prometheus
     optionalApi libs.dropwizard.metrics.core
     optionalApi libs.prometheus
+    implementation 'io.micrometer:context-propagation:1.1.1'
 
     // Netty
     api libs.netty.transport

core/src/main/java/com/linecorp/armeria/common/RequestContextAccessor.java

@@ -0,0 +1,140 @@
+/*
+ * Copyright 2020 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you 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:
+ *
+ *   https://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 com.linecorp.armeria.common;
+
+import org.reactivestreams.Subscription;
+
+import com.linecorp.armeria.internal.common.RequestContextUtil;
+
+import io.micrometer.context.ContextRegistry;
+import io.micrometer.context.ContextSnapshot;
+import io.micrometer.context.ContextSnapshot.Scope;
+import io.micrometer.context.ThreadLocalAccessor;
+
+/**
+ * This class works with the
+ * <a href="https://docs.micrometer.io/context-propagation/reference/index.html">
+ * Context-propagation</a> library and keep the {@link RequestContext} during
+ * <a href="https://github.com/reactor/reactor-core">Reactor</a> operations.
+ * Get the {@link RequestContextAccessor} to register it to the {@link ContextRegistry}.
+ * Then, {@link ContextRegistry} will use {@link RequestContextAccessor} to
+ * propagate context during the
+ * <a href="https://github.com/reactor/reactor-core">Reactor</a> operations
+ * so that you can get the context using {@link RequestContext#current()}.
+ * However, please note that you should include Mono#contextWrite(ContextView) or
+ * Flux#contextWrite(ContextView) to end of the Reactor codes.
+ * If not, {@link RequestContext} will not be keep during Reactor Operation.
+ */
+public final class RequestContextAccessor implements ThreadLocalAccessor<RequestContext> {
+
+    private static final String KEY = RequestContextAccessor.class.getName();
+
+    /**
+     * The value which obtained through {@link RequestContextAccessor},
+     * will be stored in the Context under this {@code KEY}.
+     * This method will be called by {@link ContextSnapshot} internally.
+     */
+    @Override
+    public Object key() {
+        return KEY;
+    }
+
+    /**
+     * The value which obtained through {@link RequestContextAccessor},
+     * will be stored in the Context under this {@code KEY}.
+     * User can use this method to register {@link RequestContext} to
+     * Reactor Context.
+     */
+
+    public static String accessorKey() {
+        return KEY;
+    }
+
+    /**
+     * {@link ContextSnapshot} will call this method during the execution
+     * of lambda functions in {@link ContextSnapshot#wrap(Runnable)},
+     * as well as during Mono#subscribe(), Flux#subscribe(),
+     * {@link Subscription#request(long)}, and CoreSubscriber#onSubscribe(Subscription).
+     * Following these calls, {@link ContextSnapshot#setThreadLocals()} is
+     * invoked to restore the state of {@link RequestContextStorage}.
+     * Furthermore, at the end of these methods, {@link Scope#close()} is executed
+     * to revert the {@link RequestContextStorage} to its original state.
+     */
+    @Override
+    public RequestContext getValue() {
+        return RequestContextUtil.get();
+    }
+
+    /**
+     * {@link ContextSnapshot} will call this method during the execution
+     * of lambda functions in {@link ContextSnapshot#wrap(Runnable)},
+     * as well as during Mono#subscribe(), Flux#subscribe(),
+     * {@link Subscription#request(long)}, and CoreSubscriber#onSubscribe(Subscription).
+     * Following these calls, {@link ContextSnapshot#setThreadLocals()} is
+     * invoked to restore the state of {@link RequestContextStorage}.
+     * Furthermore, at the end of these methods, {@link Scope#close()} is executed
+     * to revert the {@link RequestContextStorage} to its original state.
+     */
+    @Override
+    public void setValue(RequestContext value) {
+        RequestContextUtil.getAndSet(value);
+    }
+
+    /**
+     * This method will be called at the start of {@link ContextSnapshot.Scope} and
+     * the end of {@link ContextSnapshot.Scope}. If reactor Context does not
+     * contains {@link RequestContextAccessor#KEY}, {@link ContextSnapshot} will use
+     * this method to remove the value from {@link ThreadLocal}.
+     * Please note that {@link RequestContextUtil#pop()} return {@link AutoCloseable} instance,
+     * but it is not used in `Try with Resources` syntax. this is because {@link ContextSnapshot.Scope}
+     * will handle the {@link AutoCloseable} instance returned by {@link RequestContextUtil#pop()}.
+     */
+    @Override
+    @SuppressWarnings("MustBeClosedChecker")
+    public void setValue() {
+        RequestContextUtil.pop();
+    }
+
+    /**
+     * {@link ContextSnapshot} will call this method during the execution
+     * of lambda functions in {@link ContextSnapshot#wrap(Runnable)},
+     * as well as during Mono#subscribe(), Flux#subscribe(),
+     * {@link Subscription#request(long)}, and CoreSubscriber#onSubscribe(Subscription).
+     * Following these calls, {@link ContextSnapshot#setThreadLocals()} is
+     * invoked to restore the state of {@link RequestContextStorage}.
+     * Furthermore, at the end of these methods, {@link Scope#close()} is executed
+     * to revert the {@link RequestContextStorage} to its original state.
+     */
+    @Override
+    public void restore(RequestContext previousValue) {
+        RequestContextUtil.getAndSet(previousValue);
+    }
+
+    /**
+     * This method will be called at the start of {@link ContextSnapshot.Scope} and
+     * the end of {@link ContextSnapshot.Scope}. If reactor Context does not
+     * contains {@link RequestContextAccessor#KEY}, {@link ContextSnapshot} will use
+     * this method to remove the value from {@link ThreadLocal}.
+     * Please note that {@link RequestContextUtil#pop()} return {@link AutoCloseable} instance,
+     * but it is not used in `Try with Resources` syntax. this is because {@link ContextSnapshot.Scope}
+     * will handle the {@link AutoCloseable} instance returned by {@link RequestContextUtil#pop()}.
+     */
+    @Override
+    @SuppressWarnings("MustBeClosedChecker")
+    public void restore() {
+        RequestContextUtil.pop();
+    }
+}

core/src/main/resources/META-INF/services/io.micrometer.context.ThreadLocalAccessor

@@ -0,0 +1 @@
+com.linecorp.armeria.common.RequestContextAccessor

core/src/test/java/com/linecorp/armeria/common/RequestContextAccessorTest.java

@@ -0,0 +1,157 @@
+/*
+ * Copyright 2019 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you 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:
+ *
+ *   https://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.
+ */
+/*
+ * Copyright 2012 The Netty Project
+ *
+ * The Netty Project licenses this file to you 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 com.linecorp.armeria.common;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.List;
+
+import org.assertj.core.api.Assertions;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.internal.common.RequestContextUtil;
+
+import io.micrometer.context.ContextRegistry;
+import io.micrometer.context.ThreadLocalAccessor;
+
+class RequestContextAccessorTest {
+
+    /* Should clean up on RequestContext.
+     * because some test case does not clean up on RequestContext,
+     * and it will affect other tests if test are executed parallely.
+     */
+
+    @AfterEach
+    @SuppressWarnings("MustBeClosedChecker")
+    void cleanUp() {
+        RequestContextUtil.pop();
+    }
+
+    @Test
+    void should_be_loaded_by_SPI() {
+        final ContextRegistry ctxRegistry = ContextRegistry.getInstance();
+        final List<ThreadLocalAccessor<?>> threadLocalAccessors = ctxRegistry.getThreadLocalAccessors();
+
+        assertThat(threadLocalAccessors.size()).isGreaterThan(1);
+        assertThat(threadLocalAccessors).hasAtLeastOneElementOfType(RequestContextAccessor.class);
+    }
+
+    @Test
+    void should_return_expected_key() {
+        // Given
+        final RequestContextAccessor reqCtxAccessor = new RequestContextAccessor();
+        final String expectedValue = RequestContextAccessor.class.getName();
+
+        // When
+        final Object result = reqCtxAccessor.key();
+
+        // Then
+        assertThat(result).isEqualTo(expectedValue);
+    }
+
+    @Test
+    void should_success_set() {
+        // Given
+        final ClientRequestContext ctx = newContext();
+        final RequestContextAccessor reqCtxAccessor = new RequestContextAccessor();
+
+        // When
+        reqCtxAccessor.setValue(ctx);
+
+        // Then
+        final RequestContext currentCtx = RequestContext.current();
+        assertThat(currentCtx).isEqualTo(ctx);
+    }
+
+    @Test
+    void should_throw_NPE_when_set_null() {
+        // Given
+        final RequestContextAccessor reqCtxAccessor = new RequestContextAccessor();
+
+        // When + Then
+        Assertions.assertThatThrownBy(() -> reqCtxAccessor.setValue(null))
+                  .isInstanceOf(NullPointerException.class);
+    }
+
+    @Test
+    void should_be_null_when_setValue() {
+        // Given
+        final ClientRequestContext ctx = newContext();
+        final RequestContextAccessor reqCtxAccessor = new RequestContextAccessor();
+        reqCtxAccessor.setValue(ctx);
+
+        // When
+        reqCtxAccessor.setValue();
+
+        // Then
+        final RequestContext reqCtx = RequestContext.currentOrNull();
+        assertThat(reqCtx).isNull();
+    }
+
+    @Test
+    void should_be_restore_original_state_when_restore() {
+        // Given
+        final RequestContextAccessor reqCtxAccessor = new RequestContextAccessor();
+        final ClientRequestContext previousCtx = newContext();
+        final ClientRequestContext currentCtx = newContext();
+        reqCtxAccessor.setValue(currentCtx);
+
+        // When
+        reqCtxAccessor.restore(previousCtx);
+
+        // Then
+        final RequestContext reqCtx = RequestContext.currentOrNull();
+        assertThat(reqCtx).isNotNull();
+        assertThat(reqCtx).isEqualTo(previousCtx);
+    }
+
+    @Test
+    void should_be_null_when_restore() {
+        // Given
+        final RequestContextAccessor reqCtxAccessor = new RequestContextAccessor();
+        final ClientRequestContext currentCtx = newContext();
+        reqCtxAccessor.setValue(currentCtx);
+
+        // When
+        reqCtxAccessor.restore();
+
+        // Then
+        final RequestContext reqCtx = RequestContext.currentOrNull();
+        assertThat(reqCtx).isNull();
+    }
+
+    static ClientRequestContext newContext() {
+        return ClientRequestContext.builder(HttpRequest.of(HttpMethod.GET, "/"))
+                                   .build();
+    }
+}

reactor3/build.gradle

@@ -1,3 +1,4 @@
 dependencies {
     api libs.reactor.core
+    implementation 'io.micrometer:context-propagation:1.1.1'
 }

reactor3/src/main/java/com/linecorp/armeria/common/reactor3/RequestContextPropagationHook.java

@@ -0,0 +1,71 @@
+/*
+ * Copyright 2020 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you 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:
+ *
+ *   https://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 com.linecorp.armeria.common.reactor3;
+
+import com.linecorp.armeria.common.RequestContext;
+
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Mono;
+
+/**
+ * Utility class to keep {@link RequestContext} during
+ * <a href="https://github.com/reactor/reactor-core">Reactor</a> operations
+ * with <a href="https://docs.micrometer.io/context-propagation/reference/index.html">
+ * Context-propagation</a>.
+ */
+public final class RequestContextPropagationHook {
+
+    private static volatile boolean enabled;
+
+    /**
+     * Enable <a href="https://docs.micrometer.io/context-propagation/reference/index.html">
+     * Context-propagation</a> to keep {@link RequestContext} during
+     * Reactor operations.
+     * Please note that enable {@link RequestContextPropagationHook} at the
+     * start of the application. otherwise, {@link RequestContext} may not be keep.
+     */
+    public static synchronized void enable() {
+        if (enabled) {
+            return;
+        }
+        Hooks.enableAutomaticContextPropagation();
+        enabled = true;
+    }
+
+    /**
+     * It returns whether the {@link RequestContextPropagationHook} is enabled.
+     */
+    public static boolean isEnabled() {
+        return enabled;
+    }
+
+    /**
+     * It disable {@link RequestContextPropagationHook}. {@link RequestContext}
+     * will not be keep during both {@link Mono} and {@link Flux} Operations.
+     */
+    public static synchronized void disable() {
+        if (!enabled) {
+            return;
+        }
+
+        Hooks.disableAutomaticContextPropagation();
+        enabled = false;
+    }
+
+    private RequestContextPropagationHook() {}
+}