@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();
}
}
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;
}