Skip to content

Improve context propagation for brave integration #6139

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ikhoon merged 15 commits into from
Jun 5, 2025
Merged

Improve context propagation for brave integration #6139

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
5eafb19
minimal impl
jrhee17 Mar 4, 2025
6f5a662
cleanup
jrhee17 Mar 5, 2025
f2f0d90
Merge branch 'main' into feat/brave-prop
jrhee17 Mar 25, 2025
a525b8a
Merge branch 'main' into feat/brave-prop
jrhee17 Apr 3, 2025
23eeedb
add propagation logic to executors
jrhee17 Apr 4, 2025
46833c4
address comments by @minwoox and @trustin
jrhee17 Apr 10, 2025
933912a
minor nits
jrhee17 Apr 10, 2025
f737fea
warn when a different context is set
jrhee17 Apr 10, 2025
4598cc8
use object equality for checking trace context equality
jrhee17 Apr 10, 2025
2ab56cb
address comment by @minwoox
jrhee17 Apr 15, 2025
3296dcb
Merge remote-tracking branch 'origin/main' into feat/brave-prop
jrhee17 May 20, 2025
835d38d
lint
jrhee17 May 20, 2025
39ecbe8
Merge branch 'main' into feat/brave-prop
jrhee17 Jun 2, 2025
b7dee51
address comment by @ikhoon
jrhee17 Jun 2, 2025
fafd081
Merge branch 'main' into feat/brave-prop
jrhee17 Jun 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions brave/brave6/src/main/java/com/linecorp/armeria/client/brave/BraveClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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());

@minwoox minwoox Apr 8, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Because a span is fixed to a RequestContext, how about leaving a warning log if a RequestContext has a TraceContext already? I've seen internal customers who have used this API.
Also, I think we can remove @Nullable from setTraceContext method.

@jrhee17 jrhee17 Apr 10, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm unsure if leaving a log level with WARN level is a good idea since it is a valid behavior to do so. e.g. Using TraceContextPropagation#inject (span1) with BraveClient (span2) will also overwrite the span where span2.parent == span1.

What do you think of leaving a trace level log so that we can guide users if there is unexpected behavior?

Alternatively, I can also add a static field guard so that the log is left only once if you feel strongly of this.

@minwoox minwoox Apr 10, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, that's valid situtation. IIUC, in that case, the same TraceContext can be set to a RequestContext. Is it right? If so, can't we add a warning log if it tries to set a different TraceContext?

@jrhee17 jrhee17 Apr 10, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've seen internal customers who have used this API.

It would help if you could explain why users using this API is a bad idea. If users were setting the TraceContext manually before this change, couldn't they continue using it this way after this change as well?

If so, can't we add a warning log if it tries to set a different TraceContext?

Just to be clear are you suggesting that a warning log is printed once guarded by a static variable? or every time a different context is set?

@jrhee17 jrhee17 Apr 10, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm unsure if leaving a log level with WARN level is a good idea since it is a valid behavior to do so. e.g. Using TraceContextPropagation#inject (span1) with BraveClient (span2) will also overwrite the span where span2.parent == span1.

Just to be clear, the following scenario yields different spans:

HttpTracing httpTracing = HttpTracing.newBuilder(tracing).build();
ClientRequestContext ctx = ClientRequestContext.of(HttpRequest.of(HttpMethod.GET, "/"));
final HttpClientRequest braveReq = ClientRequestContextAdapter.asHttpClientRequest(ctx, RequestHeaders.builder());
ScopedSpan span1 = tracing.tracer().startScopedSpan("span1");
TraceContextUtil.setTraceContext(ctx, span1.context());
Span span2 = HttpClientHandler.create(httpTracing).handleSend(braveReq);
assert span1.context() != span2.context();

@minwoox minwoox Apr 10, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It would help if you could explain why users using this API is a bad idea. If users were setting the TraceContext manually before this change, couldn't they continue using it this way after this change as well?

I've imagined some of this situation. Please let me know if this doesn't make sense:

ClientRequextContext ctx = ...
TraceContextUtil.setTraceContext(ctx, span1.context());
ctx.makeContextAware(future).thenApply(
  // This handler will be executed later after `TraceContextUtil.setTraceContext(ctx, span2.context());` is invoked.
  // Expected to execute with span1 but executed with span2.
)

// Other sets different TraceContext to the ctx:
TraceContextUtil.setTraceContext(ctx, span2.context());

@jrhee17 jrhee17 Apr 14, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I still maintain that

  1. the scenario you shared seems like a possible case prior to the changes in this PR
  2. the warning log may print for valid scenarios mentioned in here Improve context propagation for brave integration #6139 (comment)

Anyways, I've updated to always leave a warning log as you requested. PTAL when you have time

@minwoox minwoox Apr 14, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

the warning log may print for valid scenarios mentioned in here

Since it's a valid scenario, could you change to log just one time?

@ikhoon ikhoon May 21, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for chiming in late. prevTraceContext could be the server's context where the client is used. The logging could cause confusion despite correct usage.

sb.decorator(LoggingService.newDecorator());
sb.decorator(BraveService.newDecorator(tracing));
sb.service("/", (ctx, req) -> {
    final WebClient client = WebClient.builder(server1.httpUri())
                                      .decorator(LoggingClient.newDecorator())
                                      .decorator(BraveClient.newDecorator(tracing))
                                      .build();
    return client.get("/");
});
}

Should we use ctx.ownAttr(TRACE_CONTEXT_KEY) to check the duplicate context?

@jrhee17 jrhee17 Jun 2, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, done


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);
Expand Down Expand Up @@ -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";
}
};
}
181 changes: 25 additions & 156 deletions ...e6/src/main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,17 +16,8 @@

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;
Expand Down Expand Up @@ -92,65 +83,37 @@ public static RequestContextCurrentTraceContextBuilder builder() {
* > }
* > });
* }</pre>
*
* @deprecated this setting has no effect
Comment thread
trustin marked this conversation as resolved.
Outdated
*/
@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<TraceContext> THREAD_LOCAL_CONTEXT = new ThreadLocal<>();

private static final ThreadLocal<Boolean> 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<Pattern> nonRequestThreadPatterns;
private final boolean scopeDecoratorAdded;

RequestContextCurrentTraceContext(CurrentTraceContext.Builder builder,
List<Pattern> 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 = 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
Expand All @@ -160,16 +123,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
Expand All @@ -190,104 +159,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<RequestContext> {

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;
}
}
}
12 changes: 5 additions & 7 deletions ...main/java/com/linecorp/armeria/common/brave/RequestContextCurrentTraceContextBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -31,8 +29,6 @@
*/
public final class RequestContextCurrentTraceContextBuilder extends CurrentTraceContext.Builder {

private final ImmutableList.Builder<Pattern> nonRequestThreadPatterns = ImmutableList.builder();

private boolean scopeDecoratorAdded;

RequestContextCurrentTraceContextBuilder() {}
Expand All @@ -44,7 +40,9 @@ public final class RequestContextCurrentTraceContextBuilder extends CurrentTrace
* monitoring requests.
*
* @see RequestContextCurrentTraceContext#setCurrentThreadNotRequestThread(boolean)
* @deprecated this setting has no effect
Comment thread
trustin marked this conversation as resolved.
Outdated
*/
@Deprecated
public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern) {
requireNonNull(pattern, "pattern");
final Pattern compiled = Pattern.compile(pattern);
Expand All @@ -58,9 +56,10 @@ public RequestContextCurrentTraceContextBuilder nonRequestThread(String pattern)
* RMI to serve monitoring requests.
*
* @see RequestContextCurrentTraceContext#setCurrentThreadNotRequestThread(boolean)
* @deprecated this setting has no effect
Comment thread
trustin marked this conversation as resolved.
Outdated
*/
@Deprecated
public RequestContextCurrentTraceContextBuilder nonRequestThread(Pattern pattern) {
nonRequestThreadPatterns.add(requireNonNull(pattern, "pattern"));
return this;
}

Expand All @@ -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);
}
}
2 changes: 2 additions & 0 deletions brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down
Loading