diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/client/brave/BraveClient.java b/brave/brave6/src/main/java/com/linecorp/armeria/client/brave/BraveClient.java index 798708b3aae..2c2317a9b8f 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/client/brave/BraveClient.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/client/brave/BraveClient.java @@ -35,6 +35,7 @@ import com.linecorp.armeria.common.logging.ClientConnectionTimings; import com.linecorp.armeria.internal.common.RequestContextExtension; import com.linecorp.armeria.internal.common.brave.SpanTags; +import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import brave.Span; import brave.Tracer; @@ -45,6 +46,7 @@ import brave.http.HttpClientResponse; import brave.http.HttpTracing; import brave.propagation.CurrentTraceContext; +import brave.propagation.CurrentTraceContext.Scope; /** * Decorates an {@link HttpClient} to trace outbound {@link HttpRequest}s using @@ -120,11 +122,13 @@ public HttpResponse execute(ClientRequestContext ctx, HttpRequest req) throws Ex final Span span = handler.handleSend(braveReq); req = req.withHeaders(newHeaders); ctx.updateRequest(req); + TraceContextUtil.setTraceContext(ctx, span.context()); final RequestContextExtension ctxExtension = ctx.as(RequestContextExtension.class); if (currentTraceContext != null && !span.isNoop() && ctxExtension != null) { - // Make the span the current span and run scope decorators when the ctx is pushed. - ctxExtension.hook(() -> currentTraceContext.newScope(span.context())); + // Run the scope decorators when the ctx is pushed to the thread local. + ctxExtension.hook(() -> currentTraceContext.decorateScope(span.context(), + CLIENT_REQUEST_DECORATING_SCOPE)); } maybeAddTagsToSpan(ctx, braveReq, span); @@ -191,4 +195,14 @@ private static void logTiming(Span span, String startName, String endName, long span.annotate(startTimeMicros, startName); span.annotate(startTimeMicros + TimeUnit.NANOSECONDS.toMicros(durationNanos), endName); } + + private static final Scope CLIENT_REQUEST_DECORATING_SCOPE = new Scope() { + @Override + public void close() {} + + @Override + public String toString() { + return "ClientRequestDecoratingScope"; + } + }; } diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContext.java b/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContext.java index 7029bb413a4..d741b0e7f63 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContext.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContext.java @@ -16,21 +16,13 @@ package com.linecorp.armeria.common.brave; -import static com.linecorp.armeria.internal.common.brave.TraceContextUtil.setTraceContext; import static com.linecorp.armeria.internal.common.brave.TraceContextUtil.traceContext; -import java.util.List; -import java.util.function.Function; -import java.util.function.Supplier; -import java.util.regex.Pattern; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import com.linecorp.armeria.client.brave.BraveClient; import com.linecorp.armeria.common.RequestContext; import com.linecorp.armeria.common.annotation.Nullable; import com.linecorp.armeria.common.annotation.UnstableApi; +import com.linecorp.armeria.internal.common.brave.InternalTraceContextUtil; import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import com.linecorp.armeria.server.brave.BraveService; @@ -92,65 +84,45 @@ public static RequestContextCurrentTraceContextBuilder builder() { * > } * > }); * } + * + * @deprecated This setting has no effect. */ + @Deprecated public static void setCurrentThreadNotRequestThread(boolean value) { - if (value) { - THREAD_NOT_REQUEST_THREAD.set(true); - } else { - THREAD_NOT_REQUEST_THREAD.remove(); - } } private static final RequestContextCurrentTraceContext DEFAULT = builder().build(); - private static final Logger logger = LoggerFactory.getLogger(RequestContextCurrentTraceContext.class); - - // Thread-local for storing TraceContext when invoking callbacks off the request thread. - private static final ThreadLocal THREAD_LOCAL_CONTEXT = new ThreadLocal<>(); - - private static final ThreadLocal THREAD_NOT_REQUEST_THREAD = new ThreadLocal<>(); - - private static final Scope INITIAL_REQUEST_SCOPE = new Scope() { + private static final Scope NOOP_SCOPE = new Scope() { @Override public void close() { - // Don't remove the outer-most context (client or server request) } @Override public String toString() { - return "InitialRequestScope"; + return "ArmeriaNoopScope"; } }; - private final List nonRequestThreadPatterns; private final boolean scopeDecoratorAdded; - RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, - List nonRequestThreadPatterns, boolean scopeDecoratorAdded) { + RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, boolean scopeDecoratorAdded) { super(builder); - - this.nonRequestThreadPatterns = nonRequestThreadPatterns; this.scopeDecoratorAdded = scopeDecoratorAdded; } @Override @Nullable public TraceContext get() { - final RequestContext ctx = getRequestContextOrWarnOnce(); - if (ctx == null) { - return THREAD_LOCAL_CONTEXT.get(); + final TraceContext traceContext = InternalTraceContextUtil.get(); + if (traceContext != null) { + return traceContext; } - - if (ctx.eventLoop().inEventLoop()) { - return traceContext(ctx); - } else { - final TraceContext threadLocalContext = THREAD_LOCAL_CONTEXT.get(); - if (threadLocalContext != null) { - return threadLocalContext; - } - // First span on a non-request thread will use the request's TraceContext as a parent. - return traceContext(ctx); + final RequestContext ctx = RequestContext.currentOrNull(); + if (ctx == null) { + return null; } + return traceContext(ctx); } @Override @@ -160,16 +132,27 @@ public Scope newScope(@Nullable TraceContext currentSpan) { return Scope.NOOP; } - final RequestContext ctx = getRequestContextOrWarnOnce(); + final TraceContext threadPrev = InternalTraceContextUtil.get(); + if (threadPrev == currentSpan) { + // a custom noop scope is used to avoid special behavior in built-in scope decorators + return decorateScope(currentSpan, NOOP_SCOPE); + } + + InternalTraceContextUtil.set(currentSpan); + + class ThreadLocalContextScope implements Scope { + @Override + public void close() { + InternalTraceContextUtil.set(threadPrev); + } - if (ctx != null && ctx.eventLoop().inEventLoop()) { - return createScopeForRequestThread(ctx, currentSpan); - } else { - // The RequestContext is the canonical thread-local storage for the thread processing the request. - // However, when creating spans on other threads (e.g., a thread-pool), we must use separate - // thread-local storage to prevent threads from replacing the same trace context. - return createScopeForNonRequestThread(currentSpan); + @Override + public String toString() { + return "ThreadLocalScope"; + } } + + return decorateScope(currentSpan, new ThreadLocalContextScope()); } @UnstableApi @@ -190,104 +173,4 @@ public Scope decorateScope(@Nullable TraceContext context, Scope scope) { public boolean scopeDecoratorAdded() { return scopeDecoratorAdded; } - - private Scope createScopeForRequestThread(RequestContext ctx, @Nullable TraceContext currentSpan) { - final TraceContext previous = traceContext(ctx); - setTraceContext(ctx, currentSpan); - - // Don't remove the outer-most context (client or server request) - if (previous == null) { - return decorateScope(currentSpan, INITIAL_REQUEST_SCOPE); - } - - // Removes sub-spans (i.e. local spans) from the current context when Brave's scope does. - // If an asynchronous sub-span, it may still complete later. - class RequestContextTraceContextScope implements Scope { - @Override - public void close() { - // re-lookup the attribute to avoid holding a reference to the request if this scope is leaked - final RequestContext ctx = getRequestContextOrWarnOnce(); - if (ctx != null) { - setTraceContext(ctx, previous); - } - } - - @Override - public String toString() { - return "RequestContextTraceContextScope"; - } - } - - return decorateScope(currentSpan, new RequestContextTraceContextScope()); - } - - private Scope createScopeForNonRequestThread(@Nullable TraceContext currentSpan) { - final TraceContext previous = THREAD_LOCAL_CONTEXT.get(); - THREAD_LOCAL_CONTEXT.set(currentSpan); - class ThreadLocalScope implements Scope { - @Override - public void close() { - THREAD_LOCAL_CONTEXT.set(previous); - } - - @Override - public String toString() { - return "ThreadLocalScope"; - } - } - - return decorateScope(currentSpan, new ThreadLocalScope()); - } - - /** - * Armeria code should always have a request context available, and this won't work without it. - */ - @Nullable - private RequestContext getRequestContextOrWarnOnce() { - if (Boolean.TRUE.equals(THREAD_NOT_REQUEST_THREAD.get())) { - return null; - } - if (!nonRequestThreadPatterns.isEmpty()) { - final String threadName = Thread.currentThread().getName(); - for (Pattern pattern : nonRequestThreadPatterns) { - if (pattern.matcher(threadName).find()) { - // A matched thread will match forever, so it's worth avoiding this regex match on every - // time the thread is used by saving into the ThreadLocal. - setCurrentThreadNotRequestThread(true); - return null; - } - } - } - return RequestContext.mapCurrent(Function.identity(), LogRequestContextWarningOnce.INSTANCE); - } - - private enum LogRequestContextWarningOnce implements Supplier { - - INSTANCE; - - @Override - @Nullable - public RequestContext get() { - ClassLoaderHack.loadMe(); - return null; - } - - /** - * This won't be referenced until {@link #get()} is called. If there's only one classloader, the - * initializer will only be called once. - */ - private static final class ClassLoaderHack { - static void loadMe() {} - - static { - logger.warn("Attempted to propagate trace context, but no request context available. " + - "Did you forget to use RequestContext.makeContextAware()?", - new NoRequestContextException()); - } - } - - private static final class NoRequestContextException extends RuntimeException { - private static final long serialVersionUID = 2804189311774982052L; - } - } } diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextBuilder.java b/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextBuilder.java index d3573575091..3cc19f3587e 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextBuilder.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextBuilder.java @@ -20,8 +20,6 @@ import java.util.regex.Pattern; -import com.google.common.collect.ImmutableList; - import brave.propagation.CurrentTraceContext; import brave.propagation.CurrentTraceContext.Builder; import brave.propagation.CurrentTraceContext.ScopeDecorator; @@ -31,8 +29,6 @@ */ public final class RequestContextCurrentTraceContextBuilder extends CurrentTraceContext.Builder { - private final ImmutableList.Builder nonRequestThreadPatterns = ImmutableList.builder(); - private boolean scopeDecoratorAdded; RequestContextCurrentTraceContextBuilder() {} @@ -44,7 +40,9 @@ public final class RequestContextCurrentTraceContextBuilder extends CurrentTrace * monitoring requests. * * @see RequestContextCurrentTraceContext#setCurrentThreadNotRequestThread(boolean) + * @deprecated This setting has no effect. */ + @Deprecated public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern) { requireNonNull(pattern, "pattern"); final Pattern compiled = Pattern.compile(pattern); @@ -58,9 +56,10 @@ public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern) * RMI to serve monitoring requests. * * @see RequestContextCurrentTraceContext#setCurrentThreadNotRequestThread(boolean) + * @deprecated This setting has no effect. */ + @Deprecated public RequestContextCurrentTraceContextBuilder nonRequestThread(Pattern pattern) { - nonRequestThreadPatterns.add(requireNonNull(pattern, "pattern")); return this; } @@ -79,7 +78,6 @@ public Builder addScopeDecorator(ScopeDecorator scopeDecorator) { */ @Override public RequestContextCurrentTraceContext build() { - return new RequestContextCurrentTraceContext(this, nonRequestThreadPatterns.build(), - scopeDecoratorAdded); + return new RequestContextCurrentTraceContext(this, scopeDecoratorAdded); } } diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/BraveThreadLocalAccessorProvider.java b/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/BraveThreadLocalAccessorProvider.java new file mode 100644 index 00000000000..7c4866b4b9b --- /dev/null +++ b/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/BraveThreadLocalAccessorProvider.java @@ -0,0 +1,52 @@ +/* + * Copyright 2025 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.internal.common.brave; + +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider; + +import brave.propagation.TraceContext; +import io.micrometer.context.ThreadLocalAccessor; + +public final class BraveThreadLocalAccessorProvider implements ThreadLocalAccessorProvider { + + @Override + public ThreadLocalAccessor threadLocalAccessor() { + return new ThreadLocalAccessor() { + @Override + public Object key() { + return BraveThreadLocalAccessorProvider.class; + } + + @Nullable + @Override + public Object getValue() { + return InternalTraceContextUtil.get(); + } + + @Override + public void setValue(Object o) { + InternalTraceContextUtil.set((TraceContext) o); + } + + @Override + public void setValue() { + InternalTraceContextUtil.set(null); + } + }; + } +} diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/InternalTraceContextUtil.java b/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/InternalTraceContextUtil.java new file mode 100644 index 00000000000..57bd319ccbc --- /dev/null +++ b/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/InternalTraceContextUtil.java @@ -0,0 +1,43 @@ +/* + * Copyright 2025 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.internal.common.brave; + +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext; + +import brave.propagation.TraceContext; + +/** + * Internal class for manipulate the internal {@link ThreadLocal} for {@link RequestContextCurrentTraceContext}. + * This class is reserved for internal usage and is subject to behavior change any time. + */ +public final class InternalTraceContextUtil { + + // Thread-local for storing TraceContext when invoking callbacks off the request thread. + private static final ThreadLocal THREAD_LOCAL_CONTEXT = new ThreadLocal<>(); + + public static void set(@Nullable TraceContext traceContext) { + THREAD_LOCAL_CONTEXT.set(traceContext); + } + + @Nullable + public static TraceContext get() { + return THREAD_LOCAL_CONTEXT.get(); + } + + private InternalTraceContextUtil() {} +} diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/TraceContextUtil.java b/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/TraceContextUtil.java index ad3f5890b52..34c2532e32a 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/TraceContextUtil.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/TraceContextUtil.java @@ -18,6 +18,11 @@ import static java.util.Objects.requireNonNull; +import java.util.Objects; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import com.google.common.annotations.VisibleForTesting; import com.linecorp.armeria.common.RequestContext; @@ -31,15 +36,26 @@ public final class TraceContextUtil { + private static final Logger logger = LoggerFactory.getLogger(TraceContextUtil.class); + private static final AttributeKey TRACE_CONTEXT_KEY = AttributeKey.valueOf(TraceContextUtil.class, "TRACE_CONTEXT"); + private static boolean logTraceContextOverwrite; @Nullable public static TraceContext traceContext(RequestContext ctx) { return ctx.attr(TRACE_CONTEXT_KEY); } - public static void setTraceContext(RequestContext ctx, @Nullable TraceContext traceContext) { + public static void setTraceContext(RequestContext ctx, TraceContext traceContext) { + if (!logTraceContextOverwrite) { + final TraceContext prevTraceContext = ctx.ownAttr(TRACE_CONTEXT_KEY); + if (prevTraceContext != null && !Objects.equals(traceContext, prevTraceContext)) { + logTraceContextOverwrite = true; + logger.warn("Overriding traceContext<{}> with new traceContext<{}> for RequestContext<{}>.", + prevTraceContext, traceContext, ctx.id()); + } + } ctx.setAttr(TRACE_CONTEXT_KEY, traceContext); } diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/AbstractBraveService.java b/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/AbstractBraveService.java index 20f9197ec72..4ff056eb790 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/AbstractBraveService.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/AbstractBraveService.java @@ -24,6 +24,7 @@ import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext; import com.linecorp.armeria.common.logging.RequestLog; import com.linecorp.armeria.internal.common.RequestContextExtension; +import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import com.linecorp.armeria.server.Service; import com.linecorp.armeria.server.ServiceRequestContext; import com.linecorp.armeria.server.SimpleDecoratingService; @@ -56,6 +57,7 @@ public final O serve(ServiceRequestContext ctx, I req) throws Exception { } final BI braveReq = braveRequest(ctx); final Span span = handleReceive(braveReq); + TraceContextUtil.setTraceContext(ctx, span.context()); final RequestContextExtension ctxExtension = ctx.as(RequestContextExtension.class); if (currentTraceContext.scopeDecoratorAdded() && !span.isNoop() && ctxExtension != null) { diff --git a/brave/brave6/src/main/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider b/brave/brave6/src/main/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider new file mode 100644 index 00000000000..44d87416812 --- /dev/null +++ b/brave/brave6/src/main/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider @@ -0,0 +1 @@ +com.linecorp.armeria.internal.common.brave.BraveThreadLocalAccessorProvider diff --git a/brave/brave6/src/test/java/com/linecorp/armeria/client/brave/TraceContextPropagationTest.java b/brave/brave6/src/test/java/com/linecorp/armeria/client/brave/TraceContextPropagationTest.java index 42ad4c5f77d..f07468599e3 100644 --- a/brave/brave6/src/test/java/com/linecorp/armeria/client/brave/TraceContextPropagationTest.java +++ b/brave/brave6/src/test/java/com/linecorp/armeria/client/brave/TraceContextPropagationTest.java @@ -23,21 +23,27 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.RegisterExtension; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.CsvSource; import com.linecorp.armeria.client.BlockingWebClient; +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.client.ClientRequestContextCaptor; +import com.linecorp.armeria.client.Clients; import com.linecorp.armeria.client.WebClient; import com.linecorp.armeria.common.AggregatedHttpResponse; import com.linecorp.armeria.common.HttpResponse; import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext; +import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import com.linecorp.armeria.server.ServerBuilder; import com.linecorp.armeria.server.ServiceRequestContext; import com.linecorp.armeria.server.annotation.Blocking; import com.linecorp.armeria.server.annotation.Get; import com.linecorp.armeria.server.brave.BraveService; import com.linecorp.armeria.server.logging.LoggingService; +import com.linecorp.armeria.testing.junit5.common.EventLoopExtension; import com.linecorp.armeria.testing.junit5.server.ServerExtension; import brave.Tracing; @@ -56,6 +62,10 @@ class TraceContextPropagationTest { .build(); private static final Map traceContexts = new ConcurrentHashMap<>(); + private static final String RESCHEDULE_HEADER = "x-reschedule"; + + @RegisterExtension + static EventLoopExtension eventLoop = new EventLoopExtension(); @RegisterExtension static ServerExtension server = new ServerExtension() { @@ -131,6 +141,22 @@ public HttpResponse foo() { }); sb.decorator(BraveService.newDecorator(tracing)); sb.decorator(LoggingService.newDecorator()); + + sb.decorator((delegate, ctx, req) -> { + if (req.headers().contains(RESCHEDULE_HEADER)) { + final CompletableFuture cf = new CompletableFuture<>(); + eventLoop.get().execute(() -> { + try { + cf.complete(delegate.serve(ctx, req)); + } catch (Throwable t) { + cf.completeExceptionally(t); + } + }); + return HttpResponse.of(cf); + } + return delegate.serve(ctx, req); + }); + sb.service("/hello", (ctx, req) -> HttpResponse.of(currentTraceContext.get().traceIdString())); } }; @@ -164,4 +190,17 @@ void propagation(String type) { assertThat(barServiceTraceContext.parentId()).isEqualTo(fooServiceTraceContext.spanId()); assertThat(barServiceTraceContext.spanId()).isEqualTo(barClientTraceContext.spanId()); } + + @Test + void rescheduleBeforeBraveService() { + try (ClientRequestContextCaptor captor = Clients.newContextCaptor()) { + final AggregatedHttpResponse res = + server.webClient(cb -> cb.addHeader(RESCHEDULE_HEADER, true) + .decorator(BraveClient.newDecorator(tracing))) + .blocking().get("/hello"); + assertThat(res.status().code()).isEqualTo(200); + final ClientRequestContext cctx = captor.get(); + assertThat(res.contentUtf8()).isEqualTo(TraceContextUtil.traceContext(cctx).traceIdString()); + } + } } diff --git a/brave/brave6/src/test/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextTest.java b/brave/brave6/src/test/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextTest.java index 9b7e5ed5713..03498efd4fe 100644 --- a/brave/brave6/src/test/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextTest.java +++ b/brave/brave6/src/test/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextTest.java @@ -17,14 +17,19 @@ package com.linecorp.armeria.common.brave; import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; import static org.mockito.Mockito.when; +import java.util.concurrent.atomic.AtomicReference; + import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import com.linecorp.armeria.client.ClientRequestContext; import com.linecorp.armeria.common.HttpMethod; import com.linecorp.armeria.common.HttpRequest; import com.linecorp.armeria.common.RequestContext; @@ -32,6 +37,7 @@ import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import com.linecorp.armeria.internal.common.brave.TraceContextUtil.PingPongExtra; import com.linecorp.armeria.server.ServiceRequestContext; +import com.linecorp.armeria.testing.junit5.common.EventLoopExtension; import brave.propagation.CurrentTraceContext; import brave.propagation.CurrentTraceContext.Scope; @@ -49,6 +55,9 @@ class RequestContextCurrentTraceContextTest { RequestContextCurrentTraceContext.ofDefault(); private static final TraceContext traceContext = TraceContext.newBuilder().traceId(1).spanId(1).build(); + @RegisterExtension + static final EventLoopExtension eventLoopExtension = new EventLoopExtension(); + @BeforeEach void setUp() { when(eventLoop.inEventLoop()).thenReturn(true); @@ -83,7 +92,6 @@ public void newScope_appliesWhenNoCurrentRequestContext() { public void newScope_appliesWhenCurrentRequestContext() { try (SafeCloseable requestContextScope = ctx.push()) { try (Scope traceContextScope = currentTraceContext.newScope(traceContext)) { - assertThat(traceContextScope).hasToString("InitialRequestScope"); assertThat(currentTraceContext.get()).isEqualTo(traceContext); } } @@ -92,14 +100,13 @@ public void newScope_appliesWhenCurrentRequestContext() { @Test public void newScope_closeDoesntClearFirstScope() { final TraceContext traceContext2 = TraceContext.newBuilder().traceId(1).spanId(2).build(); + TraceContextUtil.setTraceContext(ctx, traceContext); try (SafeCloseable requestContextScope = ctx.push()) { try (Scope traceContextScope = currentTraceContext.newScope(traceContext)) { - assertThat(traceContextScope).hasToString("InitialRequestScope"); assertThat(currentTraceContext.get()).isEqualTo(traceContext); try (Scope traceContextScope2 = currentTraceContext.newScope(traceContext2)) { - assertThat(traceContextScope2).hasToString("RequestContextTraceContextScope"); assertThat(currentTraceContext.get()).isEqualTo(traceContext2); } assertThat(currentTraceContext.get()).isEqualTo(traceContext); @@ -112,10 +119,10 @@ public void newScope_closeDoesntClearFirstScope() { @Test public void newScope_notOnEventLoop() { final TraceContext traceContext2 = TraceContext.newBuilder().traceId(1).spanId(2).build(); + TraceContextUtil.setTraceContext(ctx, traceContext); try (SafeCloseable requestContextScope = ctx.push()) { try (Scope traceContextScope = currentTraceContext.newScope(traceContext)) { - assertThat(traceContextScope).hasToString("InitialRequestScope"); assertThat(currentTraceContext.get()).isEqualTo(traceContext); when(eventLoop.inEventLoop()).thenReturn(false); @@ -166,4 +173,34 @@ public void shouldSetPongIfOnlyExtra() { assertThat(extra.isPong()).isTrue(); } + + @Test + void nonEventLoopPropagation() { + final AtomicReference ref = new AtomicReference<>(); + final RequestContext ctx = ClientRequestContext.of(HttpRequest.of(HttpMethod.GET, "/")); + try (Scope scope = currentTraceContext.newScope(traceContext)) { + ctx.makeContextAware(eventLoopExtension.get()).submit(() -> { + final TraceContext traceContext = currentTraceContext.get(); + ref.set(traceContext); + }); + } + await().untilAsserted(() -> assertThat(ref).doesNotHaveNullValue()); + assertThat(ref.get()).isSameAs(traceContext); + } + + @Test + void nonEventLoopFallback() { + final AtomicReference ref = new AtomicReference<>(); + final RequestContext ctx = ClientRequestContext.of(HttpRequest.of(HttpMethod.GET, "/")); + // fallback to the last context set to the RequestContext manually + TraceContextUtil.setTraceContext(ctx, traceContext); + + ctx.makeContextAware(eventLoopExtension.get()).submit(() -> { + final TraceContext traceContext = currentTraceContext.get(); + ref.set(traceContext); + }); + + await().untilAsserted(() -> assertThat(ref).doesNotHaveNullValue()); + assertThat(ref.get()).isSameAs(traceContext); + } } diff --git a/brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java b/brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java index 358e663b0c7..6f150c32a41 100644 --- a/brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java +++ b/brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java @@ -246,7 +246,6 @@ private static TestService.AsyncIface newClient(String path) { private static Tracing newTracing(String name) { final CurrentTraceContext currentTraceContext = RequestContextCurrentTraceContext.builder() - .nonRequestThread("nonrequest-") .addScopeDecorator(StrictScopeDecorator.create()) .build(); return Tracing.newBuilder() @@ -540,21 +539,16 @@ private static void testClientTimesOut(TestService.Iface client) { @Test void testNoRequestContextTraceable() throws Exception { - RequestContextCurrentTraceContext.setCurrentThreadNotRequestThread(true); - try { - final Tracing tracing = newTracing("no-request"); - final ScopedSpan span1 = tracing.tracer().startScopedSpan("span1"); - final ScopedSpan span2 = tracing.tracer().startScopedSpan("span2"); + final Tracing tracing = newTracing("no-request"); + final ScopedSpan span1 = tracing.tracer().startScopedSpan("span1"); + final ScopedSpan span2 = tracing.tracer().startScopedSpan("span2"); - assertThat(span2.context().traceId()).isEqualTo(span1.context().traceId()); + assertThat(span2.context().traceId()).isEqualTo(span1.context().traceId()); - span2.finish(); - span1.finish(); + span2.finish(); + span1.finish(); - spanHandler.take(2); - } finally { - RequestContextCurrentTraceContext.setCurrentThreadNotRequestThread(false); - } + spanHandler.take(2); } @Test diff --git a/core/build.gradle b/core/build.gradle index 054381732c2..612b0947c99 100644 --- a/core/build.gradle +++ b/core/build.gradle @@ -107,6 +107,7 @@ dependencies { // Micrometer and other metric-related stuff api libs.micrometer.core + api libs.context.propagation optionalApi libs.micrometer.prometheus.legacy optionalApi libs.dropwizard.metrics.core optionalApi libs.prometheus.legacy diff --git a/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutor.java b/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutor.java index af0fcdd0bd9..16b165e347b 100644 --- a/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutor.java +++ b/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutor.java @@ -24,6 +24,7 @@ import org.slf4j.LoggerFactory; import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; abstract class AbstractContextAwareExecutor implements Executor { enum LogRequestContextWarningOnce implements Supplier { @@ -73,7 +74,10 @@ public final E withoutContext() { final Runnable makeContextAware(Runnable task) { final RequestContext context = contextOrNull(); - return context == null ? task : context.makeContextAware(task); + if (context != null) { + return context.makeContextAware(task); + } + return ArmeriaContextPropagation.captureAll().wrap(task); } @Override diff --git a/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutorService.java b/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutorService.java index 3f8211001a8..a47cc43e1fd 100644 --- a/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutorService.java +++ b/core/src/main/java/com/linecorp/armeria/common/AbstractContextAwareExecutorService.java @@ -27,6 +27,8 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + abstract class AbstractContextAwareExecutorService extends AbstractContextAwareExecutor implements ExecutorService { @@ -101,7 +103,10 @@ public final T invokeAny( final Callable makeContextAware(Callable task) { final RequestContext context = contextOrNull(); - return context == null ? task : context.makeContextAware(task); + if (context != null) { + return context.makeContextAware(task); + } + return ArmeriaContextPropagation.captureAll().wrap(task); } private Collection> makeContextAware( diff --git a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiConsumer.java b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiConsumer.java index 00a9dce93b6..f58932a0bf7 100644 --- a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiConsumer.java +++ b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiConsumer.java @@ -21,14 +21,20 @@ import java.util.function.BiConsumer; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshot.Scope; final class DefaultContextAwareBiConsumer implements ContextAwareBiConsumer { private final RequestContext context; private final BiConsumer action; + private final ContextSnapshot contextSnapshot; DefaultContextAwareBiConsumer(RequestContext context, BiConsumer action) { this.context = requireNonNull(context, "context"); this.action = requireNonNull(action, "action"); + contextSnapshot = ArmeriaContextPropagation.captureAll(); } @Override @@ -44,7 +50,9 @@ public BiConsumer withoutContext() { @Override public void accept(T t, U u) { try (SafeCloseable ignored = context.push()) { - action.accept(t, u); + try (Scope ignored2 = contextSnapshot.setThreadLocals()) { + action.accept(t, u); + } } } } diff --git a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiFunction.java b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiFunction.java index 10ec70e28e1..b1b66bc4d59 100644 --- a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiFunction.java +++ b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareBiFunction.java @@ -21,15 +21,21 @@ import java.util.function.BiFunction; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshot.Scope; final class DefaultContextAwareBiFunction implements ContextAwareBiFunction { private final RequestContext context; private final BiFunction function; + private final ContextSnapshot contextSnapshot; DefaultContextAwareBiFunction(RequestContext context, BiFunction function) { this.context = requireNonNull(context, "context"); this.function = requireNonNull(function, "function"); + contextSnapshot = ArmeriaContextPropagation.captureAll(); } @Override @@ -45,7 +51,9 @@ public BiFunction withoutContext() { @Override public R apply(T t, U u) { try (SafeCloseable ignored = context.push()) { - return function.apply(t, u); + try (Scope ignored2 = contextSnapshot.setThreadLocals()) { + return function.apply(t, u); + } } } } diff --git a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareCallable.java b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareCallable.java index 86b8f5f619e..8cee4747e46 100644 --- a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareCallable.java +++ b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareCallable.java @@ -21,14 +21,20 @@ import java.util.concurrent.Callable; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshot.Scope; final class DefaultContextAwareCallable implements ContextAwareCallable { private final RequestContext context; private final Callable callable; + private final ContextSnapshot contextSnapshot; DefaultContextAwareCallable(RequestContext context, Callable callable) { this.context = requireNonNull(context, "context"); this.callable = requireNonNull(callable, "callable"); + contextSnapshot = ArmeriaContextPropagation.captureAll(); } @Override @@ -44,7 +50,9 @@ public Callable withoutContext() { @Override public T call() throws Exception { try (SafeCloseable ignored = context.push()) { - return callable.call(); + try (Scope ignored2 = contextSnapshot.setThreadLocals()) { + return callable.call(); + } } } } diff --git a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareConsumer.java b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareConsumer.java index b980f664b4c..b44c872cff4 100644 --- a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareConsumer.java +++ b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareConsumer.java @@ -21,14 +21,20 @@ import java.util.function.Consumer; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshot.Scope; final class DefaultContextAwareConsumer implements ContextAwareConsumer { private final RequestContext context; private final Consumer action; + private final ContextSnapshot contextSnapshot; DefaultContextAwareConsumer(RequestContext context, Consumer action) { this.context = requireNonNull(context, "context"); this.action = requireNonNull(action, "action"); + contextSnapshot = ArmeriaContextPropagation.captureAll(); } @Override @@ -44,7 +50,9 @@ public Consumer withoutContext() { @Override public void accept(T t) { try (SafeCloseable ignored = context.push()) { - action.accept(t); + try (Scope ignored2 = contextSnapshot.setThreadLocals()) { + action.accept(t); + } } } } diff --git a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareFunction.java b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareFunction.java index 054d556c8b2..4e5e73cf648 100644 --- a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareFunction.java +++ b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareFunction.java @@ -21,15 +21,21 @@ import java.util.function.Function; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshot.Scope; final class DefaultContextAwareFunction implements ContextAwareFunction { private final RequestContext context; private final Function function; + private final ContextSnapshot contextSnapshot; DefaultContextAwareFunction(RequestContext context, Function function) { this.context = requireNonNull(context, "context"); this.function = requireNonNull(function, "function"); + contextSnapshot = ArmeriaContextPropagation.captureAll(); } @Override @@ -45,7 +51,9 @@ public Function withoutContext() { @Override public R apply(T t) { try (SafeCloseable ignored = context.push()) { - return function.apply(t); + try (Scope ignored2 = contextSnapshot.setThreadLocals()) { + return function.apply(t); + } } } } diff --git a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareRunnable.java b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareRunnable.java index d1dc7993fbf..b2ceca93d9f 100644 --- a/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareRunnable.java +++ b/core/src/main/java/com/linecorp/armeria/common/DefaultContextAwareRunnable.java @@ -19,14 +19,20 @@ import static java.util.Objects.requireNonNull; import com.linecorp.armeria.common.util.SafeCloseable; +import com.linecorp.armeria.internal.common.context.ArmeriaContextPropagation; + +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshot.Scope; final class DefaultContextAwareRunnable implements ContextAwareRunnable { private final RequestContext context; private final Runnable runnable; + private final ContextSnapshot contextSnapshot; DefaultContextAwareRunnable(RequestContext context, Runnable runnable) { this.context = requireNonNull(context, "context"); this.runnable = requireNonNull(runnable, "runnable"); + contextSnapshot = ArmeriaContextPropagation.captureAll(); } @Override @@ -42,7 +48,9 @@ public Runnable withoutContext() { @Override public void run() { try (SafeCloseable ignored = context.push()) { - runnable.run(); + try (Scope ignored2 = contextSnapshot.setThreadLocals()) { + runnable.run(); + } } } } diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/context/ArmeriaContextPropagation.java b/core/src/main/java/com/linecorp/armeria/internal/common/context/ArmeriaContextPropagation.java new file mode 100644 index 00000000000..68ede3af14d --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/internal/common/context/ArmeriaContextPropagation.java @@ -0,0 +1,47 @@ +/* + * Copyright 2025 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.internal.common.context; + +import java.util.ServiceLoader; + +import io.micrometer.context.ContextRegistry; +import io.micrometer.context.ContextSnapshot; +import io.micrometer.context.ContextSnapshotFactory; + +public final class ArmeriaContextPropagation { + + private static final ContextRegistry REGISTRY = new ContextRegistry(); + + static { + // A separate registry is used to avoid pushing the same context twice when using + // a different framework (e.g. reactor) with armeria + final ServiceLoader loader = + ServiceLoader.load(ThreadLocalAccessorProvider.class); + loader.forEach(provider -> REGISTRY.registerThreadLocalAccessor(provider.threadLocalAccessor())); + } + + private static final ContextSnapshotFactory globalFactory = ContextSnapshotFactory + .builder() + .contextRegistry(REGISTRY) + .build(); + + public static ContextSnapshot captureAll() { + return globalFactory.captureAll(); + } + + private ArmeriaContextPropagation() {} +} diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/context/ThreadLocalAccessorProvider.java b/core/src/main/java/com/linecorp/armeria/internal/common/context/ThreadLocalAccessorProvider.java new file mode 100644 index 00000000000..7d60ad7dd44 --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/internal/common/context/ThreadLocalAccessorProvider.java @@ -0,0 +1,24 @@ +/* + * Copyright 2025 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.internal.common.context; + +import io.micrometer.context.ThreadLocalAccessor; + +public interface ThreadLocalAccessorProvider { + + ThreadLocalAccessor threadLocalAccessor(); +} diff --git a/core/src/main/java/com/linecorp/armeria/internal/common/context/package-info.java b/core/src/main/java/com/linecorp/armeria/internal/common/context/package-info.java new file mode 100644 index 00000000000..0c10dde167c --- /dev/null +++ b/core/src/main/java/com/linecorp/armeria/internal/common/context/package-info.java @@ -0,0 +1,23 @@ +/* + * Copyright 2021 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. + */ + +/** + * Various classes used internally. Anything in this package can be changed or removed at any time. + */ +@NonNullByDefault +package com.linecorp.armeria.internal.common.context; + +import com.linecorp.armeria.common.annotation.NonNullByDefault; diff --git a/it/context-propagation/build.gradle b/it/context-propagation/build.gradle new file mode 100644 index 00000000000..e69de29bb2d diff --git a/it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/ContextPropagationFlagTest.java b/it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/ContextPropagationFlagTest.java new file mode 100644 index 00000000000..773dc54a1f6 --- /dev/null +++ b/it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/ContextPropagationFlagTest.java @@ -0,0 +1,54 @@ +/* + * Copyright 2022 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.internal.common; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.awaitility.Awaitility.await; + +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.RegisterExtension; + +import com.linecorp.armeria.client.ClientRequestContext; +import com.linecorp.armeria.common.HttpMethod; +import com.linecorp.armeria.common.HttpRequest; +import com.linecorp.armeria.common.RequestContext; +import com.linecorp.armeria.internal.common.MyThreadLocalAccessorProvider.MyThreadLocalAccessor; +import com.linecorp.armeria.testing.junit5.common.EventLoopExtension; + +class ContextPropagationFlagTest { + + @RegisterExtension + static final EventLoopExtension eventLoop = new EventLoopExtension(); + + @Test + void basicCase() throws InterruptedException { + final RequestContext ctx = ClientRequestContext.of(HttpRequest.of(HttpMethod.GET, "/")); + final AtomicReference atomicRef = new AtomicReference<>(); + final String armeria = "armeria"; + + MyThreadLocalAccessor.THREAD_LOCAL.set(armeria); + + ctx.makeContextAware(eventLoop.get()).execute(() -> { + atomicRef.set(MyThreadLocalAccessor.THREAD_LOCAL.get()); + }); + + await().untilAsserted(() -> assertThat(atomicRef).doesNotHaveNullValue()); + assertThat(atomicRef).hasValue(armeria); + } +} diff --git a/it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/MyThreadLocalAccessorProvider.java b/it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/MyThreadLocalAccessorProvider.java new file mode 100644 index 00000000000..807412f8baa --- /dev/null +++ b/it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/MyThreadLocalAccessorProvider.java @@ -0,0 +1,55 @@ +/* + * Copyright 2025 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.internal.common; + +import com.linecorp.armeria.common.annotation.Nullable; +import com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider; + +import io.micrometer.context.ThreadLocalAccessor; + +public final class MyThreadLocalAccessorProvider implements ThreadLocalAccessorProvider { + + static final class MyThreadLocalAccessor implements ThreadLocalAccessor { + + public static final ThreadLocal THREAD_LOCAL = new ThreadLocal<>(); + + @Override + public Object key() { + return MyThreadLocalAccessor.class; + } + + @Override + public String getValue() { + return THREAD_LOCAL.get(); + } + + @Override + public void setValue(@Nullable String s) { + THREAD_LOCAL.set(s); + } + + @Override + public void setValue() { + THREAD_LOCAL.set(null); + } + } + + @Override + public ThreadLocalAccessor threadLocalAccessor() { + return new MyThreadLocalAccessor(); + } +} diff --git a/it/context-propagation/src/test/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider b/it/context-propagation/src/test/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider new file mode 100644 index 00000000000..b08657cd9c6 --- /dev/null +++ b/it/context-propagation/src/test/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider @@ -0,0 +1 @@ +com.linecorp.armeria.internal.common.MyThreadLocalAccessorProvider diff --git a/settings.gradle b/settings.gradle index caa278f3ad1..b2ccf57f0dc 100644 --- a/settings.gradle +++ b/settings.gradle @@ -205,6 +205,7 @@ includeWithFlags ':benchmarks:jmh', 'java', 'kotlin' includeWithFlags ':benchmarks:ghz', 'java' includeWithFlags ':it:builders', 'java' includeWithFlags ':it:context-storage', 'java' +includeWithFlags ':it:context-propagation', 'java' includeWithFlags ':it:dgs', 'java17' includeWithFlags ':it:flags-cyclic-dep', 'java' includeWithFlags ':it:flags-provider', 'java', 'relocate'