Skip to content

Commit c53cb0a

Browse files
jrhee17jrhee17
committed
cleanup
1 parent dd69a72 commit c53cb0a

File tree

1 file changed

+29
-39
lines changed

1 file changed

+29
-39
lines changed

‎brave/brave6/src/test/java/com/linecorp/armeria/it/brave/BraveIntegrationTest.java

Lines changed: 29 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static com.linecorp.armeria.common.SessionProtocol.H1C;
2323
import static org.assertj.core.api.Assertions.assertThat;
2424
import static org.assertj.core.api.Assertions.assertThatThrownBy;
25+
import static org.awaitility.Awaitility.await;
2526

2627
import java.time.Duration;
2728
import java.util.ArrayList;
@@ -30,19 +31,16 @@
3031
import java.util.List;
3132
import java.util.Map;
3233
import java.util.Objects;
33-
import java.util.concurrent.BlockingDeque;
3434
import java.util.concurrent.BlockingQueue;
3535
import java.util.concurrent.CompletableFuture;
3636
import java.util.concurrent.CountDownLatch;
3737
import java.util.concurrent.Executors;
38-
import java.util.concurrent.LinkedBlockingDeque;
3938
import java.util.concurrent.LinkedBlockingQueue;
4039
import java.util.concurrent.TimeUnit;
4140
import java.util.stream.IntStream;
4241

4342
import org.apache.thrift.async.AsyncMethodCallback;
4443
import org.apache.thrift.transport.TTransportException;
45-
import org.junit.jupiter.api.AfterAll;
4644
import org.junit.jupiter.api.AfterEach;
4745
import org.junit.jupiter.api.Test;
4846
import org.junit.jupiter.api.extension.RegisterExtension;
@@ -100,9 +98,8 @@
10098
class BraveIntegrationTest {
10199

102100
private static final String CLIENT_TYPE_HEADER = "x-client-type";
101+
private static final String TIMEOUT_HEADER = "x-timeout";
103102
private static final SpanHandlerImpl spanHandler = new SpanHandlerImpl();
104-
private static final BlockingDeque<Tracing> perTestTracings = new LinkedBlockingDeque<>();
105-
private static final BlockingDeque<Tracing> perClassTracings = new LinkedBlockingDeque<>();
106103

107104
@RegisterExtension
108105
static ServerExtension server = new ServerExtension(true) {
@@ -150,7 +147,8 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
150147
countDownLatch.countDown();
151148
countDownLatch.await();
152149
}
153-
final Span span = Tracing.currentTracer().nextSpan().start();
150+
final Span span = Tracing.currentTracer().nextSpan()
151+
.name("aloha1:" + i).start();
154152
try (SpanInScope unused =
155153
Tracing.currentTracer().withSpanInScope(span)) {
156154
if (i == 1) {
@@ -171,8 +169,9 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
171169
result -> allAsList(IntStream.range(1, 3).mapToObj(
172170
i -> executorService.submit(
173171
RequestContext.current().makeContextAware(() -> {
174-
final ScopedSpan span = Tracing.currentTracer()
175-
.startScopedSpan("aloha");
172+
final ScopedSpan span =
173+
Tracing.currentTracer()
174+
.startScopedSpan("aloha2:" + i);
176175
try {
177176
return null;
178177
} finally {
@@ -190,9 +189,15 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
190189
}
191190
}));
192191

193-
sb.service("/timeout", tHttpDecorate("service/timeout",
194-
// This service never calls the handler and will timeout.
195-
(name, resultHandler) -> {}));
192+
sb.service("/timeout",
193+
tHttpDecorate("service/timeout",
194+
// This service never calls the handler and will timeout.
195+
(name, resultHandler) -> {
196+
final ServiceRequestContext ctx = ServiceRequestContext.current();
197+
if (ctx.request().headers().contains(TIMEOUT_HEADER)) {
198+
ctx.timeoutNow();
199+
}
200+
}));
196201

197202
sb.service("/http", (req, ctx) -> HttpResponse.of(HttpStatus.OK));
198203
}
@@ -201,13 +206,6 @@ protected HttpResponse doGet(ServiceRequestContext ctx, HttpRequest req)
201206
@AfterEach
202207
void afterEach() {
203208
assertThat(spanHandler.spans).isEmpty();
204-
perTestTracings.forEach(Tracing::close);
205-
perTestTracings.clear();
206-
}
207-
208-
@AfterAll
209-
static void afterAll() {
210-
perClassTracings.forEach(Tracing::close);
211209
}
212210

213211
private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
@@ -219,14 +217,14 @@ private static HttpService tHttpDecorate(String name, AsyncIface asyncIface) {
219217
return (ctx, req) -> {
220218
final String braveServiceType = ctx.request().headers().get(CLIENT_TYPE_HEADER);
221219
if ("http".equals(braveServiceType)) {
222-
return BraveService.newDecorator(newTracing(name, false)).apply(service).serve(ctx, req);
220+
return BraveService.newDecorator(newTracing(name)).apply(service).serve(ctx, req);
223221
}
224222
return service.serve(ctx, req);
225223
};
226224
}
227225

228226
private static HttpService httpDecorate(String name, HttpService service) {
229-
return BraveService.newDecorator(newTracing(name, false)).apply(service);
227+
return BraveService.newDecorator(newTracing(name)).apply(service);
230228
}
231229

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

242240
private static Tracing newTracing(String name) {
243-
return newTracing(name, true);
244-
}
245-
246-
private static Tracing newTracing(String name, boolean perTest) {
247241
final CurrentTraceContext currentTraceContext =
248242
RequestContextCurrentTraceContext.builder()
249243
.nonRequestThread("nonrequest-")
250244
.addScopeDecorator(StrictScopeDecorator.create())
251245
.build();
252-
final Tracing tracing = Tracing.newBuilder()
253-
.currentTraceContext(currentTraceContext)
254-
.localServiceName(name)
255-
.addSpanHandler(spanHandler)
256-
.sampler(Sampler.ALWAYS_SAMPLE)
257-
.build();
258-
if (perTest) {
259-
perTestTracings.add(tracing);
260-
} else {
261-
perClassTracings.add(tracing);
262-
}
263-
return tracing;
246+
return Tracing.newBuilder()
247+
.currentTraceContext(currentTraceContext)
248+
.localServiceName(name)
249+
.addSpanHandler(spanHandler)
250+
.sampler(Sampler.ALWAYS_SAMPLE)
251+
.build();
264252
}
265253

266254
@Test
@@ -457,6 +445,7 @@ void testServiceInitiatedTrace(String type) throws Exception {
457445
@Test
458446
void testSpanInThreadPoolHasSameTraceId() throws Exception {
459447
server.webClient().get("pool").aggregate().get();
448+
await().untilAsserted(() -> assertThat(spanHandler.spans).hasSize(5));
460449
final MutableSpan[] spans = spanHandler.take(5);
461450
assertThat(Arrays.stream(spans).map(MutableSpan::traceId).collect(toImmutableSet())).hasSize(1);
462451
assertThat(Arrays.stream(spans).map(MutableSpan::parentId)
@@ -473,6 +462,7 @@ void testServerTimesOut(String type) throws Exception {
473462
.path("/timeout")
474463
.factory(cf)
475464
.addHeader(CLIENT_TYPE_HEADER, type)
465+
.addHeader(TIMEOUT_HEADER, true)
476466
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
477467
.build(TestService.Iface.class);
478468
assertThatThrownBy(() -> timeoutClient.hello("name"))
@@ -499,7 +489,7 @@ void testHttp2ClientTimesOut(String type) throws Exception {
499489
.path("/timeout")
500490
.addHeader(CLIENT_TYPE_HEADER, type)
501491
.decorator(BraveClient.newDecorator(newTracing("client/timeout")))
502-
.responseTimeout(Duration.ofSeconds(3))
492+
.responseTimeout(Duration.ofSeconds(1))
503493
.build(TestService.Iface.class);
504494
testClientTimesOut(timeoutClientClientTimesOut);
505495
}
@@ -608,8 +598,8 @@ public void onError(Exception exception) {
608598
}
609599
}
610600

611-
private static class SpanHandlerImpl extends SpanHandler {
612-
private final BlockingQueue<MutableSpan> spans = new LinkedBlockingQueue<>();
601+
static final class SpanHandlerImpl extends SpanHandler {
602+
final BlockingQueue<MutableSpan> spans = new LinkedBlockingQueue<>();
613603

614604
@Override
615605
public boolean end(TraceContext context, MutableSpan span, Cause cause) {

0 commit comments

Comments
 (0)