Java NIO（四）开发基于帧的ping-pong服务器

TCP有个特性是粘包，也就是一个包后面可能跟着下一个包的部分或者全部数据，帧格式就是用来解决粘包问题的。帧格式一般都是在包的前面增加一个包头，包头里的一个字段会保存包的大小。所以基于帧的服务器的每条连接的读取数据状态机都是在“读包头”和“读包体”之间转化。

具体服务器的代码如下：

public class FramedServer {

	final public static int DEFAULT_PORT = 30303;

	public static class Packet {
		ByteBuffer bb;

		public Packet() {
			bb = null;
			//ByteBuffer.allocate(0);
		}

		public static Packet createPacket(String s) {
			Packet p = new Packet();
			byte[] bytes = s.getBytes();
			p.bb = ByteBuffer.allocate(4+bytes.length);
			p.bb.putInt(-1);
			p.bb.put(bytes);
			p.bb.putInt(0, p.bb.limit() - 4);
			p.bb.rewind();
			return p;
		}
	}

	public static class Connection {
		SocketChannel sc;
		SelectionKey scKey;

		List<Packet> outgoingQueue = new LinkedList<Packet>();

		ByteBuffer lenBuf = ByteBuffer.allocate(4);
		ByteBuffer incomingBuf = lenBuf;
		
		boolean readingHeader = true;

		public Connection(SocketChannel sc, SelectionKey scKey) {
			this.sc = sc;
			this.scKey = scKey;
		}
	}

	ServerSocketChannel serverChannel;
	Selector selector;

	public void start() throws IOException {
		selector = Selector.open();
		serverChannel = ServerSocketChannel.open();
		serverChannel.bind(new InetSocketAddress(DEFAULT_PORT));
		serverChannel.configureBlocking(false);
		serverChannel.register(selector, SelectionKey.OP_ACCEPT);
		while (true) {
			selector.select();
			Set<SelectionKey> readyKeys = selector.selectedKeys();
			Iterator<SelectionKey> it = readyKeys.iterator();
			while (it.hasNext()) {
				SelectionKey key = it.next();
				it.remove();
				if (key.isAcceptable()) {
					ServerSocketChannel server = (ServerSocketChannel) key.channel();
					SocketChannel sc = server.accept();
					if (sc != null) {
						sc.configureBlocking(false);
						SelectionKey scKey = sc.register(selector, SelectionKey.OP_READ);
						scKey.attach(new Connection(sc, scKey));
					} else {
						System.out.println("WARN: accept null.");
					}
				} else if (key.isReadable()) {
					doRead((Connection) key.attachment());
				} else if (key.isWritable()) {
					doWrite((Connection) key.attachment());
				}
			}
		}
	}

	public void doRead(Connection conn) {
		try {
			if (conn.readingHeader) {
				//reading header
				int ret = conn.sc.read(conn.lenBuf);
				if (ret < 0) {
					conn.scKey.cancel();
					conn.sc.close();
                                        return;
				}
				System.out.println("ret="+ret);
				if (conn.lenBuf.hasRemaining()) {
					return;
				} else {
					conn.lenBuf.rewind();
					int len = conn.lenBuf.getInt();
					System.out.println("header: "+len);
					conn.incomingBuf = ByteBuffer.allocate(len);
					conn.readingHeader = false;
				}
			} else {
				//reading body
				int ret = conn.sc.read(conn.incomingBuf);
				if (ret < 0) {
					conn.scKey.cancel();
					conn.sc.close();
                                        return;
				}
				if (conn.incomingBuf.hasRemaining()) {
					return;
				} else {
					conn.incomingBuf.rewind();
					System.out.println("body: "+conn.incomingBuf);
					String s = new String(conn.incomingBuf.array(), 0, conn.incomingBuf.limit());
					System.out.println("received: " + s);
					conn.incomingBuf = null;
					conn.lenBuf.rewind();
					conn.readingHeader = true;
					
					Packet p = Packet.createPacket("pong");
					conn.outgoingQueue.add(p);
					doWrite(conn);
				}
			}
		} catch (IOException e) {
			e.printStackTrace();
			conn.scKey.cancel();
			try {
				conn.sc.close();
			} catch (IOException e1) {
				e1.printStackTrace();
			}
			return;
		}
	}

	public void doWrite(Connection conn) {
		while (conn.outgoingQueue.size() > 0) {
			try {
				Packet p = conn.outgoingQueue.get(0);
				conn.sc.write(p.bb);
				if (p.bb.hasRemaining()) {
					enableWrite(conn.scKey);
					return;
				} else {
					conn.outgoingQueue.remove(0);
				}
			} catch (IOException e) {
				e.printStackTrace();
				conn.scKey.cancel();
				try {
					conn.sc.close();
				} catch (IOException e1) {
					e1.printStackTrace();
				}
				return;
			}
		}

		disableWrite(conn.scKey);
	}

	protected void enableWrite(SelectionKey scKey) {
		int ops = scKey.interestOps();
		scKey.interestOps(ops | SelectionKey.OP_WRITE);
	}

	protected void disableWrite(SelectionKey scKey) {
		int ops = scKey.interestOps();
		scKey.interestOps(ops & (~SelectionKey.OP_WRITE));
	}

	public static void main(String args[]) throws IOException {
		FramedServer server = new FramedServer();
		server.start();
	}
}

客户端代码如下：

public class FramedClient {
	
	public void readPacket(SocketChannel sc) throws IOException {
		ByteBuffer lenBuf = ByteBuffer.allocate(4);
		sc.read(lenBuf);
		lenBuf.rewind();
		int len = lenBuf.getInt();
		ByteBuffer incomingBuf = ByteBuffer.allocate(len);
		sc.read(incomingBuf);
		System.out.println(new String(incomingBuf.array(), 0, incomingBuf.limit()));
	}
	
	public void run() throws IOException {
		SocketChannel sc = SocketChannel.open();
		sc.connect(new InetSocketAddress(FramedServer.DEFAULT_PORT));
		Packet p = Packet.createPacket("ping");
		sc.write(p.bb);
		
		readPacket(sc);
		
		try {
			Thread.sleep(3000);
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
		p.bb.rewind();
		sc.write(p.bb);
		
		readPacket(sc);
		
		sc.close();
	}
	
	public static void main(String args[]) {
		FramedClient client = new FramedClient();
		try {
			client.run();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
}