From 5eafb19ff5d8b89cd0b46e7063c3976d1ae022ce Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Tue, 4 Mar 2025 18:46:23 +0900 Subject: [PATCH 01/10] minimal impl --- .../armeria/client/brave/BraveClient.java | 18 +- .../RequestContextCurrentTraceContext.java | 163 +++--------------- ...uestContextCurrentTraceContextBuilder.java | 2 + .../armeria/server/brave/BraveService.java | 2 + .../brave/TraceContextPropagationTest.java | 39 +++++ ...RequestContextCurrentTraceContextTest.java | 6 +- 6 files changed, 85 insertions(+), 145 deletions(-) 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..087f1b751e0 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,12 +16,9 @@ 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; @@ -92,7 +89,10 @@ 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); @@ -110,47 +110,26 @@ public static void setCurrentThreadNotRequestThread(boolean value) { private static final ThreadLocal THREAD_NOT_REQUEST_THREAD = new ThreadLocal<>(); - private static final Scope INITIAL_REQUEST_SCOPE = new Scope() { - @Override - public void close() { - // Don't remove the outer-most context (client or server request) - } - - @Override - public String toString() { - return "InitialRequestScope"; - } - }; - - private final List nonRequestThreadPatterns; private final boolean scopeDecoratorAdded; RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, List nonRequestThreadPatterns, 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 = THREAD_LOCAL_CONTEXT.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 +139,22 @@ public Scope newScope(@Nullable TraceContext currentSpan) { return Scope.NOOP; } - final RequestContext ctx = getRequestContextOrWarnOnce(); + final TraceContext threadPrev = THREAD_LOCAL_CONTEXT.get(); + THREAD_LOCAL_CONTEXT.set(currentSpan); - 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); + class ThreadLocalContextScope implements Scope { + @Override + public void close() { + THREAD_LOCAL_CONTEXT.set(threadPrev); + } + + @Override + public String toString() { + return "ThreadLocalScope"; + } } + + return decorateScope(currentSpan, new ThreadLocalContextScope()); } @UnstableApi @@ -190,104 +175,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..28b3c6afb98 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 @@ -58,7 +58,9 @@ 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; diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java b/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java index 2fbd4d905d3..c991a7efa6e 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java @@ -25,6 +25,7 @@ import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext; import com.linecorp.armeria.internal.common.RequestContextExtension; import com.linecorp.armeria.internal.common.brave.SpanTags; +import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import com.linecorp.armeria.server.HttpService; import com.linecorp.armeria.server.ServiceRequestContext; import com.linecorp.armeria.server.SimpleDecoratingHttpService; @@ -99,6 +100,7 @@ public HttpResponse serve(ServiceRequestContext ctx, HttpRequest req) throws Exc final HttpServerRequest braveReq = ServiceRequestContextAdapter.asHttpServerRequest(ctx); final Span span = handler.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/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..79cf46c6fb6 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 @@ -83,7 +83,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 +91,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 +110,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); From 6f5a662b66c4b65ee91bfebcf8af8132c21c0f79 Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Wed, 5 Mar 2025 11:22:22 +0900 Subject: [PATCH 02/10] cleanup --- .../RequestContextCurrentTraceContext.java | 18 +---------------- ...uestContextCurrentTraceContextBuilder.java | 10 +++------- .../it/brave/BraveIntegrationTest.java | 20 +++++++------------ 3 files changed, 11 insertions(+), 37 deletions(-) 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 087f1b751e0..b17f448f597 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 @@ -18,12 +18,6 @@ import static com.linecorp.armeria.internal.common.brave.TraceContextUtil.traceContext; -import java.util.List; -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; @@ -94,26 +88,16 @@ public static RequestContextCurrentTraceContextBuilder builder() { */ @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 final boolean scopeDecoratorAdded; - RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, - List nonRequestThreadPatterns, boolean scopeDecoratorAdded) { + RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, boolean scopeDecoratorAdded) { super(builder); this.scopeDecoratorAdded = scopeDecoratorAdded; } 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 28b3c6afb98..7b4667fe35a 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); @@ -62,7 +60,6 @@ public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern) */ @Deprecated public RequestContextCurrentTraceContextBuilder nonRequestThread(Pattern pattern) { - nonRequestThreadPatterns.add(requireNonNull(pattern, "pattern")); return this; } @@ -81,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/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java b/brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java index 1b9efe9c341..a365aacbb9f 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 @@ -256,7 +256,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() @@ -501,21 +500,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 From 23eeedb15dce9f705f851b41f07118080c6b44da Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Fri, 4 Apr 2025 16:21:21 +0900 Subject: [PATCH 03/10] add propagation logic to executors --- .../RequestContextCurrentTraceContext.java | 12 ++-- .../BraveThreadLocalAccessorProvider.java | 52 ++++++++++++++++++ .../brave/InternalTraceContextUtil.java | 43 +++++++++++++++ ...common.context.ThreadLocalAccessorProvider | 17 ++++++ ...RequestContextCurrentTraceContextTest.java | 39 +++++++++++++ core/build.gradle | 1 + .../common/AbstractContextAwareExecutor.java | 6 +- .../AbstractContextAwareExecutorService.java | 7 ++- .../common/DefaultContextAwareBiConsumer.java | 10 +++- .../common/DefaultContextAwareBiFunction.java | 10 +++- .../common/DefaultContextAwareCallable.java | 10 +++- .../common/DefaultContextAwareConsumer.java | 10 +++- .../common/DefaultContextAwareFunction.java | 10 +++- .../common/DefaultContextAwareRunnable.java | 10 +++- .../context/ArmeriaContextPropagation.java | 47 ++++++++++++++++ .../context/ThreadLocalAccessorProvider.java | 24 ++++++++ .../internal/common/context/package-info.java | 23 ++++++++ it/context-propagation/build.gradle | 0 .../common/ContextPropagationFlagTest.java | 54 ++++++++++++++++++ .../common/MyThreadLocalAccessorProvider.java | 55 +++++++++++++++++++ ...common.context.ThreadLocalAccessorProvider | 17 ++++++ settings.gradle | 1 + 22 files changed, 443 insertions(+), 15 deletions(-) create mode 100644 brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/BraveThreadLocalAccessorProvider.java create mode 100644 brave/brave6/src/main/java/com/linecorp/armeria/internal/common/brave/InternalTraceContextUtil.java create mode 100644 brave/brave6/src/main/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider create mode 100644 core/src/main/java/com/linecorp/armeria/internal/common/context/ArmeriaContextPropagation.java create mode 100644 core/src/main/java/com/linecorp/armeria/internal/common/context/ThreadLocalAccessorProvider.java create mode 100644 core/src/main/java/com/linecorp/armeria/internal/common/context/package-info.java create mode 100644 it/context-propagation/build.gradle create mode 100644 it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/ContextPropagationFlagTest.java create mode 100644 it/context-propagation/src/test/java/com/linecorp/armeria/internal/common/MyThreadLocalAccessorProvider.java create mode 100644 it/context-propagation/src/test/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider 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 b17f448f597..824028553fe 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 @@ -22,6 +22,7 @@ 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,9 +93,6 @@ public static void setCurrentThreadNotRequestThread(boolean value) { private static final RequestContextCurrentTraceContext DEFAULT = builder().build(); - // Thread-local for storing TraceContext when invoking callbacks off the request thread. - private static final ThreadLocal THREAD_LOCAL_CONTEXT = new ThreadLocal<>(); - private final boolean scopeDecoratorAdded; RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, boolean scopeDecoratorAdded) { @@ -105,7 +103,7 @@ public static void setCurrentThreadNotRequestThread(boolean value) { @Override @Nullable public TraceContext get() { - final TraceContext traceContext = THREAD_LOCAL_CONTEXT.get(); + final TraceContext traceContext = InternalTraceContextUtil.get(); if (traceContext != null) { return traceContext; } @@ -123,13 +121,13 @@ public Scope newScope(@Nullable TraceContext currentSpan) { return Scope.NOOP; } - final TraceContext threadPrev = THREAD_LOCAL_CONTEXT.get(); - THREAD_LOCAL_CONTEXT.set(currentSpan); + final TraceContext threadPrev = InternalTraceContextUtil.get(); + InternalTraceContextUtil.set(currentSpan); class ThreadLocalContextScope implements Scope { @Override public void close() { - THREAD_LOCAL_CONTEXT.set(threadPrev); + InternalTraceContextUtil.set(threadPrev); } @Override 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/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..ac6ad4c5bf1 --- /dev/null +++ b/brave/brave6/src/main/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider @@ -0,0 +1,17 @@ +# +# 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. +# + +com.linecorp.armeria.internal.common.brave.BraveThreadLocalAccessorProvider \ No newline at end of file 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 79cf46c6fb6..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); @@ -164,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/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..4caf915c44c --- /dev/null +++ b/it/context-propagation/src/test/resources/META-INF/services/com.linecorp.armeria.internal.common.context.ThreadLocalAccessorProvider @@ -0,0 +1,17 @@ +# +# 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. +# + +com.linecorp.armeria.internal.common.MyThreadLocalAccessorProvider \ No newline at end of file diff --git a/settings.gradle b/settings.gradle index 78de60e61fe..6509bcb3b92 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' From 46833c4c7575de5e8631a4d47d813821f2cd83ab Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 10 Apr 2025 10:30:20 +0900 Subject: [PATCH 04/10] address comments by @minwoox and @trustin --- .../RequestContextCurrentTraceContext.java | 18 +++++++++++++++++- ...questContextCurrentTraceContextBuilder.java | 4 ++-- .../common/brave/TraceContextUtil.java | 12 +++++++++++- ....common.context.ThreadLocalAccessorProvider | 18 +----------------- ....common.context.ThreadLocalAccessorProvider | 16 ---------------- 5 files changed, 31 insertions(+), 37 deletions(-) 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 824028553fe..96151342b4a 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 @@ -85,7 +85,7 @@ public static RequestContextCurrentTraceContextBuilder builder() { * > }); * } * - * @deprecated this setting has no effect + * @deprecated This setting has no effect */ @Deprecated public static void setCurrentThreadNotRequestThread(boolean value) { @@ -93,6 +93,17 @@ public static void setCurrentThreadNotRequestThread(boolean value) { private static final RequestContextCurrentTraceContext DEFAULT = builder().build(); + private static final Scope NOOP_SCOPE = new Scope() { + @Override + public void close() { + } + + @Override + public String toString() { + return "ArmeriaNoopScope"; + } + }; + private final boolean scopeDecoratorAdded; RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder, boolean scopeDecoratorAdded) { @@ -122,6 +133,11 @@ public Scope newScope(@Nullable TraceContext currentSpan) { } 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 { 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 7b4667fe35a..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 @@ -40,7 +40,7 @@ public final class RequestContextCurrentTraceContextBuilder extends CurrentTrace * monitoring requests. * * @see RequestContextCurrentTraceContext#setCurrentThreadNotRequestThread(boolean) - * @deprecated this setting has no effect + * @deprecated This setting has no effect. */ @Deprecated public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern) { @@ -56,7 +56,7 @@ public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern) * RMI to serve monitoring requests. * * @see RequestContextCurrentTraceContext#setCurrentThreadNotRequestThread(boolean) - * @deprecated this setting has no effect + * @deprecated This setting has no effect. */ @Deprecated public RequestContextCurrentTraceContextBuilder nonRequestThread(Pattern pattern) { 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..2786f5cc002 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,9 @@ import static java.util.Objects.requireNonNull; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + import com.google.common.annotations.VisibleForTesting; import com.linecorp.armeria.common.RequestContext; @@ -31,6 +34,8 @@ 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"); @@ -39,7 +44,12 @@ 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 (logger.isTraceEnabled()) { + final TraceContext prevTraceContext = traceContext(ctx); + logger.trace("[{}] Overriding traceContext<{}> with new traceContext<{}>", + ctx.id(), prevTraceContext, traceContext); + } ctx.setAttr(TRACE_CONTEXT_KEY, traceContext); } 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 index ac6ad4c5bf1..44d87416812 100644 --- 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 @@ -1,17 +1 @@ -# -# 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. -# - -com.linecorp.armeria.internal.common.brave.BraveThreadLocalAccessorProvider \ No newline at end of file +com.linecorp.armeria.internal.common.brave.BraveThreadLocalAccessorProvider 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 index 4caf915c44c..50f9f16592d 100644 --- 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 @@ -1,17 +1 @@ -# -# 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. -# - com.linecorp.armeria.internal.common.MyThreadLocalAccessorProvider \ No newline at end of file From 933912aa7d71e9ca237eab1ddba39d7085ba811d Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 10 Apr 2025 11:34:11 +0900 Subject: [PATCH 05/10] minor nits --- .../armeria/common/brave/RequestContextCurrentTraceContext.java | 2 +- ....armeria.internal.common.context.ThreadLocalAccessorProvider | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 96151342b4a..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 @@ -85,7 +85,7 @@ public static RequestContextCurrentTraceContextBuilder builder() { * > }); * } * - * @deprecated This setting has no effect + * @deprecated This setting has no effect. */ @Deprecated public static void setCurrentThreadNotRequestThread(boolean value) { 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 index 50f9f16592d..b08657cd9c6 100644 --- 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 @@ -1 +1 @@ -com.linecorp.armeria.internal.common.MyThreadLocalAccessorProvider \ No newline at end of file +com.linecorp.armeria.internal.common.MyThreadLocalAccessorProvider From f737fea846273143bc5e30b4e75d36f673a0048d Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 10 Apr 2025 15:10:00 +0900 Subject: [PATCH 06/10] warn when a different context is set --- .../armeria/internal/common/brave/TraceContextUtil.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 2786f5cc002..73cc3d6ac80 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 @@ -45,10 +45,10 @@ public static TraceContext traceContext(RequestContext ctx) { } public static void setTraceContext(RequestContext ctx, TraceContext traceContext) { - if (logger.isTraceEnabled()) { - final TraceContext prevTraceContext = traceContext(ctx); - logger.trace("[{}] Overriding traceContext<{}> with new traceContext<{}>", - ctx.id(), prevTraceContext, traceContext); + final TraceContext prevTraceContext = traceContext(ctx); + if (prevTraceContext != traceContext) { + logger.warn("[{}] Overriding traceContext<{}> with new traceContext<{}>", + ctx.id(), prevTraceContext, traceContext); } ctx.setAttr(TRACE_CONTEXT_KEY, traceContext); } From 4598cc88341de6898777ecfb09dfb08d45cf9cb5 Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Thu, 10 Apr 2025 15:31:21 +0900 Subject: [PATCH 07/10] use object equality for checking trace context equality --- .../armeria/internal/common/brave/TraceContextUtil.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 73cc3d6ac80..82f424b9b19 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,8 @@ import static java.util.Objects.requireNonNull; +import java.util.Objects; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -46,7 +48,7 @@ public static TraceContext traceContext(RequestContext ctx) { public static void setTraceContext(RequestContext ctx, TraceContext traceContext) { final TraceContext prevTraceContext = traceContext(ctx); - if (prevTraceContext != traceContext) { + if (Objects.equals(traceContext, prevTraceContext)) { logger.warn("[{}] Overriding traceContext<{}> with new traceContext<{}>", ctx.id(), prevTraceContext, traceContext); } From 2ab56cb89bc4549e281609e1016b30a19e5c792d Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Tue, 15 Apr 2025 10:08:59 +0900 Subject: [PATCH 08/10] address comment by @minwoox --- .../internal/common/brave/TraceContextUtil.java | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) 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 82f424b9b19..0b2a65da75f 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 @@ -40,6 +40,7 @@ public final class TraceContextUtil { private static final AttributeKey TRACE_CONTEXT_KEY = AttributeKey.valueOf(TraceContextUtil.class, "TRACE_CONTEXT"); + private static boolean logTraceContextOverwrite; @Nullable public static TraceContext traceContext(RequestContext ctx) { @@ -47,10 +48,13 @@ public static TraceContext traceContext(RequestContext ctx) { } public static void setTraceContext(RequestContext ctx, TraceContext traceContext) { - final TraceContext prevTraceContext = traceContext(ctx); - if (Objects.equals(traceContext, prevTraceContext)) { - logger.warn("[{}] Overriding traceContext<{}> with new traceContext<{}>", - ctx.id(), prevTraceContext, traceContext); + if (!logTraceContextOverwrite) { + final TraceContext prevTraceContext = traceContext(ctx); + 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); } From 835d38d53cd70cfd1deffc25a4f2b879de6a7272 Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Tue, 20 May 2025 17:17:50 +0900 Subject: [PATCH 09/10] lint --- .../java/com/linecorp/armeria/server/brave/BraveService.java | 3 --- 1 file changed, 3 deletions(-) diff --git a/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java b/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java index bf0c17ae8ac..ac93114ee6e 100644 --- a/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java +++ b/brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java @@ -26,9 +26,6 @@ import com.linecorp.armeria.common.HttpResponse; 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.SpanTags; -import com.linecorp.armeria.internal.common.brave.TraceContextUtil; import com.linecorp.armeria.server.HttpService; import com.linecorp.armeria.server.ServiceRequestContext; From b7dee518ab67fe5587f7d2f4f58f9b1e040df884 Mon Sep 17 00:00:00 2001 From: jrhee17 Date: Mon, 2 Jun 2025 15:28:51 +0900 Subject: [PATCH 10/10] address comment by @ikhoon --- .../armeria/internal/common/brave/TraceContextUtil.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 0b2a65da75f..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 @@ -49,7 +49,7 @@ public static TraceContext traceContext(RequestContext ctx) { public static void setTraceContext(RequestContext ctx, TraceContext traceContext) { if (!logTraceContextOverwrite) { - final TraceContext prevTraceContext = traceContext(ctx); + 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<{}>.",