opentelemetry-java-instrume.../instrumentation/netty/netty-4.1/javaagent/src/test/groovy/ClientHandler.java

52 lines
1.8 KiB
Java

/*
* Copyright The OpenTelemetry Authors
* SPDX-License-Identifier: Apache-2.0
*/
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.http.FullHttpResponse;
import io.netty.handler.codec.http.HttpObject;
import io.netty.handler.codec.http.HttpResponse;
import io.netty.handler.codec.http.LastHttpContent;
import io.netty.util.AttributeKey;
import java.util.concurrent.CompletableFuture;
/*
Bridges from async Netty world to the sync world of our http client tests.
When request initiated by a test gets a response, calls a given callback and completes given
future with response's status code.
*/
public class ClientHandler extends SimpleChannelInboundHandler<HttpObject> {
private static final AttributeKey<HttpResponse> HTTP_RESPONSE =
AttributeKey.valueOf(ClientHandler.class, "http-response");
private final CompletableFuture<Integer> responseCode;
public ClientHandler(CompletableFuture<Integer> responseCode) {
this.responseCode = responseCode;
}
@Override
public void channelRead0(ChannelHandlerContext ctx, HttpObject msg) {
if (msg instanceof FullHttpResponse) {
ctx.pipeline().remove(this);
FullHttpResponse response = (FullHttpResponse) msg;
responseCode.complete(response.getStatus().code());
} else if (msg instanceof HttpResponse) {
// Headers before body have been received, store them to use when finishing the span.
ctx.channel().attr(HTTP_RESPONSE).set((HttpResponse) msg);
} else if (msg instanceof LastHttpContent) {
ctx.pipeline().remove(this);
responseCode.complete(ctx.channel().attr(HTTP_RESPONSE).get().getStatus().code());
}
}
@Override
public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
responseCode.completeExceptionally(cause);
ctx.close();
}
}