架构师训练营第七周作业 --web 压测工具

public class Counter {
	private volatile int current;
	public final int max;
	public Counter(int max) {
		this.max = max;
	}
	public int incrementCurrentAndGet() {
		return current++;
	}
}

import java.io.Closeable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.apache.commons.lang.StringUtils;
import com.google.common.base.Preconditions;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.http.DefaultHttpRequest;
import io.netty.handler.codec.http.DefaultLastHttpContent;
import io.netty.handler.codec.http.FullHttpResponse;
import io.netty.handler.codec.http.HttpClientCodec;
import io.netty.handler.codec.http.HttpHeaderNames;
import io.netty.handler.codec.http.HttpHeaderValues;
import io.netty.handler.codec.http.HttpMethod;
import io.netty.handler.codec.http.HttpObjectAggregator;
import io.netty.handler.codec.http.HttpVersion;
public class HttpBenchmarkTool implements Closeable {
	private final CountDownLatch finishLatch;
	private final ConcurrentLinkedQueue<Long> responseTimeList = new ConcurrentLinkedQueue<>();
	private final ConcurrentHashMap<Channel, Counter> channelToCounter = new ConcurrentHashMap<>();
	private final ConcurrentHashMap<Channel, Long> channelToSendMs = new ConcurrentHashMap<>();
	private final NioEventLoopGroup eventLoop = new NioEventLoopGroup();
	private final Bootstrap bootstrap;
	public final String url;
	public final int totalRequest;
	public final int concurrency;
	public final String host;
	public HttpBenchmarkTool(String host, String url, int totalRequest, int concurrency) {
		Preconditions.checkArgument(StringUtils.isNotBlank(host), "host is empty");
		Preconditions.checkArgument(StringUtils.isNotBlank(url), "url is empty");
		Preconditions.checkArgument(totalRequest > 0, "totalRequest invalid");
		Preconditions.checkArgument(concurrency > 0, "concurrency invalid");
		this.host = host;
		this.url = url;
		this.totalRequest = totalRequest;
		this.concurrency = concurrency;
		this.finishLatch = new CountDownLatch(concurrency);
		this.bootstrap = initBootstrap();
	}
	
	private Bootstrap initBootstrap() {
		Bootstrap bootstrap = new Bootstrap();
		bootstrap.group(eventLoop)
		.channel(NioSocketChannel.class)
		.option(ChannelOption.TCP_NODELAY, true)
		.handler(new ChannelInitializer<Channel>() {
			@Override
			protected void initChannel(Channel ch) throws Exception {
				ch.pipeline().addLast(new HttpClientCodec(8192, 65536, 65536))
				.addLast(new HttpObjectAggregator(100 * 65536))
				.addLast(new SimpleChannelInboundHandler<FullHttpResponse>() {
					@Override
					protected void channelRead0(ChannelHandlerContext ctx, FullHttpResponse msg) throws Exception {
//						System.out.println(msg);
//						System.out.println("received");
						responseTimeList.add(System.currentTimeMillis() - channelToSendMs.get(ctx.channel()));
						DefaultHttpRequest req = new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, url);
						req.headers().add(HttpHeaderNames.USER_AGENT, "chrome")
						.add(HttpHeaderNames.HOST, host)
						.add(HttpHeaderNames.CONNECTION, HttpHeaderValues.KEEP_ALIVE);
						if (!trySendReq(req, ctx.channel())) {
							finishLatch.countDown();
						}
					}
				});
			}
		});
		return bootstrap;
	}
	
	public void doBenchmark() throws Exception {
		System.out.println("start to connect!");
		ArrayList<ChannelFuture> channelFutureList = new ArrayList<>(concurrency);
		for (int i = 0; i < concurrency; i++) {
			channelFutureList.add(bootstrap.connect(host, 80));
		}
		ArrayList<Channel> channelList = new ArrayList<>(concurrency);
		for (ChannelFuture f : channelFutureList) {
			channelList.add(f.sync().channel());
		}
		System.out.println("all connected!");
		final int reqPerChannel = totalRequest / concurrency;
		final int left = totalRequest % concurrency;
		int i = 0;
		for (Channel ch : channelList) {
			int max = reqPerChannel;
			if (left > 0 && i++ < left) {
				max++;
			}
			channelToCounter.put(ch, new Counter(max));
		}
		//send first request
		DefaultHttpRequest req = new DefaultHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, url);
		req.headers().add(HttpHeaderNames.USER_AGENT, "chrome")
		.add(HttpHeaderNames.HOST, host)
		.add(HttpHeaderNames.CONNECTION, HttpHeaderValues.KEEP_ALIVE);
		for (final Channel ch : channelList) {
			if (!trySendReq(req, ch)) {
				finishLatch.countDown();
			}
		}
		System.out.println("waiting");
		boolean succ = finishLatch.await(10L, TimeUnit.MINUTES);
		if (!succ) {
			throw new TimeoutException();
		}
		System.out.println("finish!");
		printResult();
	}
	private void printResult() {
		ArrayList<Long> responseTimeList = new ArrayList<>(this.responseTimeList);
		System.out.println("concurrency: " + concurrency);
		System.out.println("totalRequest: " + responseTimeList.size());
		double avgTime = responseTimeList.stream()
				.mapToLong(i -> i.longValue()).average().getAsDouble();
		System.out.println("avg response time(ms): " + avgTime);
		Collections.sort(responseTimeList);
		System.out.println("90% response time(ms): " + responseTimeList.get(responseTimeList.size() * 90 / 100 - 1));
		System.out.println("95% response time(ms): " + responseTimeList.get(responseTimeList.size() * 95 / 100 - 1));
		System.out.println("99% response time(ms): " + responseTimeList.get(responseTimeList.size() * 99 / 100 - 1));
	}
	private boolean trySendReq(DefaultHttpRequest req, final Channel ch) {
		Counter counter = channelToCounter.get(ch);
		if (counter.incrementCurrentAndGet() < counter.max) {
			ch.write(req);
			ch.writeAndFlush(new DefaultLastHttpContent()).addListener(f -> {
//				System.out.println("sent");
				channelToSendMs.put(ch, System.currentTimeMillis());
			});
			return true;
		}
		return false;
	}
	@Override
	public void close() throws IOException {
		try {
			eventLoop.shutdownGracefully().sync();
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
	}
}

public class Main {
	/**
	 * @param args
	 * @throws Exception 
	 */
	public static void main(String[] args) throws Exception {
		try (HttpBenchmarkTool tool = new HttpBenchmarkTool("www.baidu.com", "/", 100, 10)) {
			tool.doBenchmark();
		}
	}
}