From ab40a5ab796a2913f2cbe9ee20a0bf5f6659a34b Mon Sep 17 00:00:00 2001
From: Trustin Lee <trustin@linecorp.com>
Date: Tue, 9 Aug 2016 17:03:58 +0900
Subject: [PATCH] Use RpcRequest.method() as Zipkin span name

Motivation:

Pre-0.21 used a Thrift method name as Zipkin span name, but 0.21 always
uses HTTP method name.

Modifications:

- Use RpcRequest.method() as Zipkin span name if available via
  RequestLog

Result:

More conprehensive and Backward-compatible behavior
---
 .../armeria/client/tracing/AbstractTracingClient.java | 11 +++++++++--
 .../server/tracing/AbstractTracingService.java        | 10 ++++++++--
 .../it/logging/HttpTracingIntegrationTest.java        |  6 ++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/src/main/java/com/linecorp/armeria/client/tracing/AbstractTracingClient.java b/src/main/java/com/linecorp/armeria/client/tracing/AbstractTracingClient.java
index eb226f32413..d6b633338ab 100644
--- a/src/main/java/com/linecorp/armeria/client/tracing/AbstractTracingClient.java
+++ b/src/main/java/com/linecorp/armeria/client/tracing/AbstractTracingClient.java
@@ -84,7 +84,7 @@ public O execute(ClientRequestContext ctx, I req) throws Exception {
 
         ctx.requestLogFuture().thenAcceptBoth(
                 res.closeFuture(),
-                (log, unused) -> clientInterceptor.closeSpan(span, createResponseAdapter(ctx, log, res)))
+                (log, unused) -> closeSpan(ctx, span, log, res))
            .exceptionally(CompletionActions::log);
 
         return res;
@@ -127,6 +127,14 @@ protected List<KeyValueAnnotation> annotations(ClientRequestContext ctx, Request
         return annotations;
     }
 
+
+    private void closeSpan(ClientRequestContext ctx, Span span, RequestLog req, O res) {
+        if (req.hasAttr(RequestLog.RPC_REQUEST)) {
+            span.setName(req.attr(RequestLog.RPC_REQUEST).get().method());
+        }
+        clientInterceptor.closeSpan(span, createResponseAdapter(ctx, req, res));
+    }
+
     protected ClientResponseAdapter createResponseAdapter(
             ClientRequestContext ctx, RequestLog req, O res) {
 
@@ -174,5 +182,4 @@ public Collection<KeyValueAnnotation> requestAnnotations() {
             return Collections.emptyList();
         }
     }
-
 }
diff --git a/src/main/java/com/linecorp/armeria/server/tracing/AbstractTracingService.java b/src/main/java/com/linecorp/armeria/server/tracing/AbstractTracingService.java
index 0e619687407..2c23a0c696d 100644
--- a/src/main/java/com/linecorp/armeria/server/tracing/AbstractTracingService.java
+++ b/src/main/java/com/linecorp/armeria/server/tracing/AbstractTracingService.java
@@ -76,7 +76,7 @@ public O serve(ServiceRequestContext ctx, I req) throws Exception {
             if (sampled) {
                 ctx.requestLogFuture().thenAcceptBoth(
                         res.closeFuture(),
-                        (log, unused) -> serverInterceptor.closeSpan(serverSpan, createResponseAdapter(ctx, log, res)))
+                        (log, unused) -> closeSpan(ctx, serverSpan, log, res))
                    .exceptionally(CompletionActions::log);
             }
             return res;
@@ -136,6 +136,13 @@ protected List<KeyValueAnnotation> annotations(
         return annotations;
     }
 
+    private void closeSpan(ServiceRequestContext ctx, ServerSpan serverSpan, RequestLog req, O res) {
+        if (req.hasAttr(RequestLog.RPC_REQUEST)) {
+            serverSpan.getSpan().setName(req.attr(RequestLog.RPC_REQUEST).get().method());
+        }
+        serverInterceptor.closeSpan(serverSpan, createResponseAdapter(ctx, req, res));
+    }
+
     protected ServerResponseAdapter createResponseAdapter(
             ServiceRequestContext ctx, RequestLog req, O res) {
         final List<KeyValueAnnotation> annotations = annotations(ctx, req, res);
@@ -171,5 +178,4 @@ public Collection<KeyValueAnnotation> requestAnnotations() {
             return Collections.emptyList();
         }
     }
-
 }
diff --git a/src/test/java/com/linecorp/armeria/it/logging/HttpTracingIntegrationTest.java b/src/test/java/com/linecorp/armeria/it/logging/HttpTracingIntegrationTest.java
index 809f7438414..145bc3a780f 100644
--- a/src/test/java/com/linecorp/armeria/it/logging/HttpTracingIntegrationTest.java
+++ b/src/test/java/com/linecorp/armeria/it/logging/HttpTracingIntegrationTest.java
@@ -144,6 +144,9 @@ public void testClientInitiatedTrace() throws Exception {
         assertThat(spans[3].getAnnotations()).allMatch(a -> "service/bar".equals(a.host.service_name));
         assertThat(spans[4].getAnnotations()).allMatch(a -> "client/qux".equals(a.host.service_name));
         assertThat(spans[5].getAnnotations()).allMatch(a -> "service/qux".equals(a.host.service_name));
+
+        // Check the span names.
+        assertThat(spans).allMatch(s -> "hello".equals(s.getName()));
     }
 
     @Test(timeout = 10000)
@@ -178,6 +181,9 @@ public void testServiceInitiatedTrace() throws Exception {
         assertThat(spans[2].getAnnotations()).allMatch(a -> "service/bar".equals(a.host.service_name));
         assertThat(spans[3].getAnnotations()).allMatch(a -> "client/qux".equals(a.host.service_name));
         assertThat(spans[4].getAnnotations()).allMatch(a -> "service/qux".equals(a.host.service_name));
+
+        // Check the span names.
+        assertThat(spans).allMatch(s -> "hello".equals(s.getName()));
     }
 
     private static class DelegatingCallback implements AsyncMethodCallback<String> {