@Override
public O serve(ServiceRequestContext ctx, I req) throws Exception {
final TraceData traceData = getTraceData(ctx, req);
final String method = req instanceof RpcRequest ? ((RpcRequest) req).method() : ctx.method();
final ServerRequestAdapter requestAdapter = new InternalServerRequestAdapter(method, traceData);
final ServerSpan serverSpan = serverInterceptor.openSpan(requestAdapter);
final boolean sampled;
if (serverSpan != null) {
ctx.onEnter(() -> serverInterceptor.setSpan(serverSpan));
ctx.onExit(serverInterceptor::clearSpan);
sampled = serverSpan.getSample();
} else {
sampled = false;
}
try {
final O res = delegate().serve(ctx, req);
if (sampled) {
ctx.requestLogFuture()
.thenAcceptBoth(
res.closeFuture(), (log, unused) -> closeSpan(ctx, serverSpan, log, res))
.exceptionally(CompletionActions::log);
}
return res;
} finally {
serverInterceptor.clearSpan();
}
}
/** Returns the server-side annotations that should be added to a Zipkin span. */
protected List<KeyValueAnnotation> annotations(ServiceRequestContext ctx, RequestLog req, O res) {
final List<KeyValueAnnotation> annotations = new ArrayList<>(5);
final StringBuilder uriBuilder = new StringBuilder();
uriBuilder.append(req.scheme().uriText());
uriBuilder.append("://");
uriBuilder.append(req.host());
uriBuilder.append(ctx.path());
if (req.method() != null) {
uriBuilder.append('#');
uriBuilder.append(req.method());
}
annotations.add(KeyValueAnnotation.create("server.uri", uriBuilder.toString()));
if (ctx.remoteAddress() != null) {
annotations.add(KeyValueAnnotation.create("server.remote", ctx.remoteAddress().toString()));
}
if (ctx.localAddress() != null) {
annotations.add(KeyValueAnnotation.create("server.local", ctx.localAddress().toString()));
}
final CompletableFuture<?> f = res.closeFuture();
if (f.isDone()) {
// Need to use a callback because CompletableFuture does not have a getter for the cause of
// failure.
// The callback will be invoked immediately because the future is done already.
f.handle(
voidFunction(
(result, cause) -> {
final String resultText = cause == null ? "success" : "failure";
annotations.add(KeyValueAnnotation.create("server.result", resultText));
if (cause != null) {
annotations.add(KeyValueAnnotation.create("server.cause", cause.toString()));
}
}))
.exceptionally(CompletionActions::log);
}
return annotations;
}
private static SpanCollectingReporter testServiceInvocation(boolean sampled) throws Exception {
SpanCollectingReporter reporter = new SpanCollectingReporter();
Brave brave =
new Brave.Builder(TEST_SERVICE)
.reporter(reporter)
.traceSampler(Sampler.create(1.0f))
.build();
@SuppressWarnings("unchecked")
Service<ThriftCall, ThriftReply> delegate = mock(Service.class);
TraceData traceData =
TraceData.builder()
.sample(sampled)
.spanId(SpanId.builder().traceId(1).spanId(2).parentId(3L).build())
.build();
TracingServiceImpl stub = new TracingServiceImpl(delegate, brave, traceData);
ThriftCall req = new ThriftCall(0, HelloService.Iface.class, "hello", "trustin");
DefaultRequestLog reqLog = new DefaultRequestLog();
reqLog.start(mock(Channel.class), SessionProtocol.H2C, "localhost", TEST_METHOD, "/");
reqLog.end();
ServiceRequestContext ctx = mock(ServiceRequestContext.class);
// AbstractTracingService prefers RpcRequest.method() to ctx.method(), so "POST" should be
// ignored.
when(ctx.method()).thenReturn("POST");
when(ctx.requestLogFuture()).thenReturn(reqLog);
ctx.onEnter(ArgumentMatchers.isA(Runnable.class));
ctx.onExit(ArgumentMatchers.isA(Runnable.class));
ThriftReply res = new ThriftReply(0, "Hello, trustin!");
when(delegate.serve(ctx, req)).thenReturn(res);
// do invoke
stub.serve(ctx, req);
verify(delegate, times(1)).serve(eq(ctx), eq(req));
return reporter;
}