Introduce BraveRpcService

brave/brave5/build.gradle

@@ -1,6 +1,7 @@
 dependencies {
     api libs.brave5
     api libs.brave5.instrumentation.http
+    api libs.brave5.instrumentation.rpc
 
     if (project.ext.targetJavaVersion >= 11) {
         testImplementation project(':thrift0.18')

brave/brave6/build.gradle

@@ -1,6 +1,7 @@
 dependencies {
     api libs.brave6
     api libs.brave6.instrumentation.http
+    api libs.brave6.instrumentation.rpc
 
     if (project.ext.targetJavaVersion >= 11) {
         testImplementation project(':thrift0.18')

brave/brave6/src/main/java/com/linecorp/armeria/server/brave/AbstractBraveService.java

@@ -0,0 +1,93 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY 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.server.brave;
+
+import static com.linecorp.armeria.server.brave.ArmeriaServerParser.annotateWireSpan;
+import static com.linecorp.armeria.server.brave.BraveService.SERVICE_REQUEST_DECORATING_SCOPE;
+
+import com.linecorp.armeria.common.Request;
+import com.linecorp.armeria.common.Response;
+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.server.Service;
+import com.linecorp.armeria.server.ServiceRequestContext;
+import com.linecorp.armeria.server.SimpleDecoratingService;
+import com.linecorp.armeria.server.TransientServiceOption;
+
+import brave.Span;
+import brave.Tracer;
+import brave.Tracer.SpanInScope;
+
+abstract class AbstractBraveService<BI extends brave.Request, BO extends brave.Response,
+        I extends Request, O extends Response> extends SimpleDecoratingService<I, O> {
+
+    private final Tracer tracer;
+    private final RequestContextCurrentTraceContext currentTraceContext;
+
+    /**
+     * Creates a new instance that decorates the specified {@link Service}.
+     */
+    protected AbstractBraveService(Service<I, O> delegate, Tracer tracer,
+                                   RequestContextCurrentTraceContext currentTraceContext) {
+        super(delegate);
+        this.tracer = tracer;
+        this.currentTraceContext = currentTraceContext;
+    }
+
+    @Override
+    public final O serve(ServiceRequestContext ctx, I req) throws Exception {
+        if (!ctx.config().transientServiceOptions().contains(TransientServiceOption.WITH_TRACING)) {
+            return unwrap().serve(ctx, req);
+        }
+        final BI braveReq = braveRequest(ctx);
+        final Span span = handleReceive(braveReq);
+
+        final RequestContextExtension ctxExtension = ctx.as(RequestContextExtension.class);
+        if (currentTraceContext.scopeDecoratorAdded() && !span.isNoop() && ctxExtension != null) {
+            // Run the scope decorators when the ctx is pushed to the thread local.
+            ctxExtension.hook(() -> currentTraceContext.decorateScope(span.context(),
+                                                                        SERVICE_REQUEST_DECORATING_SCOPE));
+        }
+
+        maybeAddTagsToSpan(ctx, braveReq, span);
+        try (SpanInScope ignored = tracer.withSpanInScope(span)) {
+            return unwrap().serve(ctx, req);
+        }
+    }
+
+    abstract BI braveRequest(ServiceRequestContext ctx);
+
+    abstract BO braveResponse(ServiceRequestContext ctx, RequestLog log, BI braveReq);
+
+    abstract Span handleReceive(BI braveReq);
+
+    abstract void handleSend(BO response, Span span);
+
+    void maybeAddTagsToSpan(ServiceRequestContext ctx, BI braveReq, Span span) {
+        if (span.isNoop()) {
+            // For no-op spans, nothing special to do.
+            return;
+        }
+
+        ctx.log().whenComplete().thenAccept(log -> {
+            annotateWireSpan(log, span);
+            final BO braveRes = braveResponse(ctx, log, braveReq);
+            handleSend(braveRes, span);
+        });
+    }
+}

brave/brave6/src/main/java/com/linecorp/armeria/server/brave/ArmeriaHttpServerParser.java → brave/brave6/src/main/java/com/linecorp/armeria/server/brave/ArmeriaServerParser.java

@@ -16,51 +16,26 @@
 
 package com.linecorp.armeria.server.brave;
 
-import com.linecorp.armeria.common.RpcRequest;
 import com.linecorp.armeria.common.logging.RequestLog;
 import com.linecorp.armeria.internal.common.brave.SpanTags;
 import com.linecorp.armeria.server.ServiceRequestContext;
 
+import brave.Request;
+import brave.Response;
+import brave.Span;
 import brave.SpanCustomizer;
-import brave.http.HttpRequestParser;
-import brave.http.HttpResponse;
-import brave.http.HttpResponseParser;
 import brave.propagation.TraceContext;
 
-/**
- * Default implementation of {@link HttpRequestParser} and {@link HttpResponseParser} for servers.
- * This parser adds some custom tags and overwrites the name of span if {@link RequestLog#requestContent()}
- * is {@link RpcRequest}.
- * The following tags become available:
- * <ul>
- *   <li>http.url</li>
- *   <li>http.host</li>
- *   <li>http.protocol</li>
- *   <li>http.serfmt</li>
- *   <li>address.remote</li>
- *   <li>address.local</li>
- * </ul>
- */
-final class ArmeriaHttpServerParser implements HttpRequestParser, HttpResponseParser {
-
-    private static final ArmeriaHttpServerParser INSTANCE = new ArmeriaHttpServerParser();
-
-    static ArmeriaHttpServerParser get() {
-        return INSTANCE;
-    }
+final class ArmeriaServerParser {
 
-    private ArmeriaHttpServerParser() {
+    private ArmeriaServerParser() {
     }
 
-    @Override
-    public void parse(brave.http.HttpRequest request, TraceContext context, SpanCustomizer span) {
-        HttpRequestParser.DEFAULT.parse(request, context, span);
-
-        final Object unwrapped = request.unwrap();
+    static void parseRequest(Request req, TraceContext context, SpanCustomizer span) {
+        final Object unwrapped = req.unwrap();
         if (!(unwrapped instanceof ServiceRequestContext)) {
             return;
         }
-
         final ServiceRequestContext ctx = (ServiceRequestContext) unwrapped;
         span.tag(SpanTags.TAG_HTTP_HOST, ctx.request().authority())
             .tag(SpanTags.TAG_HTTP_URL, ctx.request().uri().toString())
@@ -69,16 +44,12 @@ public void parse(brave.http.HttpRequest request, TraceContext context, SpanCust
             .tag(SpanTags.TAG_ADDRESS_LOCAL, ctx.localAddress().toString());
     }
 
-    @Override
-    public void parse(HttpResponse response, TraceContext context, SpanCustomizer span) {
-        HttpResponseParser.DEFAULT.parse(response, context, span);
-
-        final Object res = response.unwrap();
-        if (!(res instanceof ServiceRequestContext)) {
+    static void parseResponse(Response res, TraceContext context, SpanCustomizer span) {
+        final Object unwrapped = res.unwrap();
+        if (!(unwrapped instanceof ServiceRequestContext)) {
             return;
         }
-
-        final ServiceRequestContext ctx = (ServiceRequestContext) res;
+        final ServiceRequestContext ctx = (ServiceRequestContext) unwrapped;
         final RequestLog requestLog = ctx.log().ensureComplete();
         final String serFmt = ServiceRequestContextAdapter.serializationFormat(requestLog);
         if (serFmt != null) {
@@ -90,4 +61,18 @@ public void parse(HttpResponse response, TraceContext context, SpanCustomizer sp
             span.name(name);
         }
     }
+
+    static void annotateWireSpan(RequestLog log, Span span) {
+        span.start(log.requestStartTimeMicros());
+        final Long wireReceiveTimeNanos = log.requestFirstBytesTransferredTimeNanos();
+        assert wireReceiveTimeNanos != null;
+        SpanTags.logWireReceive(span, wireReceiveTimeNanos, log);
+
+        final Long wireSendTimeNanos = log.responseFirstBytesTransferredTimeNanos();
+        if (wireSendTimeNanos != null) {
+            SpanTags.logWireSend(span, wireSendTimeNanos, log);
+        } else {
+            // If the client timed-out the request, we will have never sent any response data at all.
+        }
+    }
 }

brave/brave6/src/main/java/com/linecorp/armeria/server/brave/BraveRpcService.java

@@ -0,0 +1,105 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY 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.server.brave;
+
+import static com.linecorp.armeria.internal.common.brave.TraceContextUtil.ensureScopeUsesRequestContext;
+
+import java.util.function.Function;
+
+import com.linecorp.armeria.common.RpcRequest;
+import com.linecorp.armeria.common.RpcResponse;
+import com.linecorp.armeria.common.annotation.UnstableApi;
+import com.linecorp.armeria.common.brave.RequestContextCurrentTraceContext;
+import com.linecorp.armeria.common.logging.RequestLog;
+import com.linecorp.armeria.server.RpcService;
+import com.linecorp.armeria.server.ServiceRequestContext;
+
+import brave.Span;
+import brave.Tracing;
+import brave.rpc.RpcRequestParser;
+import brave.rpc.RpcResponseParser;
+import brave.rpc.RpcServerHandler;
+import brave.rpc.RpcServerRequest;
+import brave.rpc.RpcServerResponse;
+import brave.rpc.RpcTracing;
+
+/**
+ * Decorates an {@link RpcService} to trace inbound {@link RpcRequest}s using
+ * <a href="https://github.com/openzipkin/brave">Brave</a>.
+ */
+@UnstableApi
+public final class BraveRpcService extends AbstractBraveService<RpcServerRequest, RpcServerResponse,
+        RpcRequest, RpcResponse> implements RpcService {
+
+    private static final RpcRequestParser defaultRequestParser = (request, context, span) -> {
+        RpcRequestParser.DEFAULT.parse(request, context, span);
+        BraveServerParsers.rpcRequestParser().parse(request, context, span);
+    };
+
+    private static final RpcResponseParser defaultResponseParser = (response, context, span) -> {
+        RpcResponseParser.DEFAULT.parse(response, context, span);
+        BraveServerParsers.rpcResponseParser().parse(response, context, span);
+    };
+
+    /**
+     * Creates a new tracing {@link RpcService} decorator using the specified {@link Tracing} instance.
+     */
+    public static Function<? super RpcService, BraveRpcService>
+    newDecorator(Tracing tracing) {
+        return newDecorator(RpcTracing.newBuilder(tracing)
+                                      .serverRequestParser(defaultRequestParser)
+                                      .serverResponseParser(defaultResponseParser)
+                                      .build());
+    }
+
+    /**
+     * Creates a new tracing {@link RpcService} decorator using the specified {@link RpcTracing} instance.
+     */
+    public static Function<? super RpcService, BraveRpcService>
+    newDecorator(RpcTracing rpcTracing) {
+        ensureScopeUsesRequestContext(rpcTracing.tracing());
+        return service -> new BraveRpcService(service, rpcTracing);
+    }
+
+    private final RpcServerHandler handler;
+
+    private BraveRpcService(RpcService delegate, RpcTracing rpcTracing) {
+        super(delegate, rpcTracing.tracing().tracer(),
+              (RequestContextCurrentTraceContext) rpcTracing.tracing().currentTraceContext());
+        handler = RpcServerHandler.create(rpcTracing);
+    }
+
+    @Override
+    RpcServerRequest braveRequest(ServiceRequestContext ctx) {
+        return RpcServiceRequestContextAdapter.asRpcServerRequest(ctx);
+    }
+
+    @Override
+    RpcServerResponse braveResponse(ServiceRequestContext ctx, RequestLog log, RpcServerRequest braveReq) {
+        return RpcServiceRequestContextAdapter.asRpcServerResponse(ctx, log, braveReq);
+    }
+
+    @Override
+    Span handleReceive(RpcServerRequest braveReq) {
+        return handler.handleReceive(braveReq);
+    }
+
+    @Override
+    void handleSend(RpcServerResponse response, Span span) {
+        handler.handleSend(response, span);
+    }
+}