Skip to content

Commit

Permalink
cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
@jrhee17
jrhee17 committed Feb 21, 2025
1 parent dd69a72 commit c53cb0a
Showing 1 changed file with 29 additions and 39 deletions.
68 changes: 29 additions & 39 deletions brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import static com.linecorp.armeria.common.SessionProtocol.H1C;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.awaitility.Awaitility.await;

import java.time.Duration;
import java.util.ArrayList;
Expand All @@ -30,19 +31,16 @@
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.BlockingDeque;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.stream.IntStream;

import org.apache.thrift.async.AsyncMethodCallback;
import org.apache.thrift.transport.TTransportException;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;
Expand Down Expand Up @@ -100,9 +98,8 @@
class BraveIntegrationTest {

private static final String CLIENT_TYPE_HEADER = "x-client-type";
private static final String TIMEOUT_HEADER = "x-timeout";
private static final SpanHandlerImpl spanHandler = new SpanHandlerImpl();
private static final BlockingDeque<Tracing> perTestTracings = new LinkedBlockingDeque<>();
private static final BlockingDeque<Tracing> perClassTracings = new LinkedBlockingDeque<>();

@RegisterExtension
static ServerExtension server = new ServerExtension(true) {
Expand Down Expand Up @@ -150,7 +147,8 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
countDownLatch.countDown();
countDownLatch.await();
}
final Span span = Tracing.currentTracer().nextSpan().start();
final Span span = Tracing.currentTracer().nextSpan()
.name("aloha1:" + i).start();
try (SpanInScope unused =
Tracing.currentTracer().withSpanInScope(span)) {
if (i == 1) {
Expand All @@ -171,8 +169,9 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
result -> allAsList(IntStream.range(1, 3).mapToObj(
i -> executorService.submit(
RequestContext.current().makeContextAware(() -> {
final ScopedSpan span = Tracing.currentTracer()
.startScopedSpan("aloha");
final ScopedSpan span =
Tracing.currentTracer()
.startScopedSpan("aloha2:" + i);
try {
return null;
} finally {
Expand All @@ -190,9 +189,15 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
}
}));

sb.service("/timeout", tHttpDecorate("service/timeout",
// This service never calls the handler and will timeout.
(name, resultHandler) -> {}));
sb.service("/timeout",
tHttpDecorate("service/timeout",
// This service never calls the handler and will timeout.
(name, resultHandler) -> {
final ServiceRequestContext ctx = ServiceRequestContext.current();
if (ctx.request().headers().contains(TIMEOUT_HEADER)) {
ctx.timeoutNow();
}
}));

sb.service("/http", (req, ctx) -> HttpResponse.of(HttpStatus.OK));
}
Expand All @@ -201,13 +206,6 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
@AfterEach
void afterEach() {
assertThat(spanHandler.spans).isEmpty();
perTestTracings.forEach(Tracing::close);
perTestTracings.clear();
}

@AfterAll
static void afterAll() {
perClassTracings.forEach(Tracing::close);
}

private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
Expand All @@ -219,14 +217,14 @@ private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
return (ctx, req) -> {
final String braveServiceType = ctx.request().headers().get(CLIENT_TYPE_HEADER);
if ("http".equals(braveServiceType)) {
return BraveService.newDecorator(newTracing(name, false)).apply(service).serve(ctx, req);
return BraveService.newDecorator(newTracing(name)).apply(service).serve(ctx, req);
}
return service.serve(ctx, req);
};
}

private static HttpService httpDecorate(String name, HttpService service) {
return BraveService.newDecorator(newTracing(name, false)).apply(service);
return BraveService.newDecorator(newTracing(name)).apply(service);
}

private static TestService.AsyncIface newClient(String path) {
Expand All @@ -240,27 +238,17 @@ private static TestService.AsyncIface newClient(String path) {
}

private static Tracing newTracing(String name) {
return newTracing(name, true);
}

private static Tracing newTracing(String name, boolean perTest) {
final CurrentTraceContext currentTraceContext =
RequestContextCurrentTraceContext.builder()
.nonRequestThread("nonrequest-")
.addScopeDecorator(StrictScopeDecorator.create())
.build();
final Tracing tracing = Tracing.newBuilder()
.currentTraceContext(currentTraceContext)
.localServiceName(name)
.addSpanHandler(spanHandler)
.sampler(Sampler.ALWAYS_SAMPLE)
.build();
if (perTest) {
perTestTracings.add(tracing);
} else {
perClassTracings.add(tracing);
}
return tracing;
return Tracing.newBuilder()
.currentTraceContext(currentTraceContext)
.localServiceName(name)
.addSpanHandler(spanHandler)
.sampler(Sampler.ALWAYS_SAMPLE)
.build();
}

@Test
Expand Down Expand Up @@ -457,6 +445,7 @@ void testServiceInitiatedTrace(String type) throws Exception {
@Test
void testSpanInThreadPoolHasSameTraceId() throws Exception {
server.webClient().get("pool").aggregate().get();
await().untilAsserted(() -> assertThat(spanHandler.spans).hasSize(5));
final MutableSpan[] spans = spanHandler.take(5);
assertThat(Arrays.stream(spans).map(MutableSpan::traceId).collect(toImmutableSet())).hasSize(1);
assertThat(Arrays.stream(spans).map(MutableSpan::parentId)
Expand All @@ -473,6 +462,7 @@ void testServerTimesOut(String type) throws Exception {
.path("/timeout")
.factory(cf)
.addHeader(CLIENT_TYPE_HEADER, type)
.addHeader(TIMEOUT_HEADER, true)
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
.build(TestService.Iface.class);
assertThatThrownBy(() -> timeoutClient.hello("name"))
Expand All @@ -499,7 +489,7 @@ void testHttp2ClientTimesOut(String type) throws Exception {
.path("/timeout")
.addHeader(CLIENT_TYPE_HEADER, type)
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
.responseTimeout(Duration.ofSeconds(3))
.responseTimeout(Duration.ofSeconds(1))
.build(TestService.Iface.class);
testClientTimesOut(timeoutClientClientTimesOut);
}
Expand Down Expand Up @@ -608,8 +598,8 @@ public void onError(Exception exception) {
}
}

private static class SpanHandlerImpl extends SpanHandler {
private final BlockingQueue<MutableSpan> spans = new LinkedBlockingQueue<>();
static final class SpanHandlerImpl extends SpanHandler {
final BlockingQueue<MutableSpan> spans = new LinkedBlockingQueue<>();

@Override
public boolean end(TraceContext context, MutableSpan span, Cause cause) {
Expand Down

0 comments on commit c53cb0a

Please sign in to comment.