一個關於NIO的分析

最近給內部做了一個NIO的分享，是基於JKD1.6的JDK的，由於我不喜歡寫PPT，所以就只寫了一個DEMO，現在把代碼拿出來分享一下，關於NIO的使用方法，以及如何擴展都在代碼的註釋裏面寫着的，希望對需要的同學有幫助。

import java.io.IOException;
import java.io.UnsupportedEncodingException;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Date;
import java.util.Iterator;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
* @author coffee mail:[email protected]
*/
public class NIoTest {

private static Logger logger = LoggerFactory.getLogger(NIoTest.class);

private Selector acceptSelector;
private Selector rwSelector;
private BlockingQueue<SocketChannel> waitRegeditChannel = new LinkedBlockingQueue<SocketChannel>();

public static void main(String[] args) {
NIoTest ns = new NIoTest();
ns.start();
}

public void start() {
InetSocketAddress localAddress = new InetSocketAddress("127.0.0.1", 8888);
ServerSocketChannel serverSocketChannel;
try {
acceptSelector = Selector.open();
rwSelector = Selector.open();
serverSocketChannel = ServerSocketChannel.open();
// 非堵塞
serverSocketChannel.configureBlocking(false);
ServerSocket socket = serverSocketChannel.socket();
// 端口不復用
socket.setReuseAddress(false);
socket.setSoTimeout(60000);
socket.setReceiveBufferSize(1024);
socket.bind(localAddress);
serverSocketChannel.register(acceptSelector, SelectionKey.OP_ACCEPT);
Executor e = Executors.newFixedThreadPool(2);// 這裏可以不用線程池
e.execute(new Accept());
e.execute(new RWThread());
} catch (IOException e) {
e.printStackTrace();
}
}

public class Accept implements Runnable {
@Override
public void run() {
while (true) {
try {
int count = acceptSelector.select(500);
// logger.debug("accept");
if (count > 0) {
Iterator<SelectionKey> keys = acceptSelector.selectedKeys().iterator();
while (keys.hasNext()) {
SelectionKey key = keys.next();
// 一定要刪除
keys.remove();
ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key.channel();
// 接受了才能獲取連接的通道
SocketChannel socketChannel = serverSocketChannel.accept();
socketChannel.configureBlocking(false);
// 取消以下注釋代碼，會導致通道在選擇器中註冊的時候與選擇器在選擇的時候爭搶互斥鎖，很難被註冊進去。
// logger.debug("開始註冊連接");
// socketChannel.register(rwSelector,
// SelectionKey.OP_READ);
// logger.debug("結束註冊連接");
waitRegeditChannel.put(socketChannel);
// 當然，可以建立一個選擇器池，併發處理接受的連接，具體如何實現自己擴展
rwSelector.wakeup();
}
}
} catch (Exception e) {
e.printStackTrace();
}
}
}
}

private class RWThread implements Runnable {
/*
* (non-Javadoc)
*
* @see java.lang.Thread#run()
*/
@Override
public void run() {
while (true) {
try {
while (!waitRegeditChannel.isEmpty()) {
SocketChannel socketChannel = waitRegeditChannel.poll();
socketChannel.register(rwSelector, SelectionKey.OP_READ);// 此處需要改造
logger.debug("註冊了一個連接:" + socketChannel.socket());
}
int count = rwSelector.select(1000);
// logger.debug("rw");
if (count > 0) {
Iterator<SelectionKey> keys = rwSelector.selectedKeys().iterator();
while (keys.hasNext()) {
SelectionKey key = keys.next();
keys.remove();
// 此處可以擴展爲將數據放到線程池中處理，這樣可以提高數據的吞吐量，但是要注意內存的處理
processKey(key);
}
}
} catch (Exception e) {
e.printStackTrace();
}
}
}

private void processKey(SelectionKey key) {
SocketChannel socketChannel = (SocketChannel) key.channel();
ByteBuffer bb = ByteBuffer.allocate(1024);
int count;
try {
// 此處加斷點以後可以明顯看到，OS底層的TCP會緩存數據，read的時候將會一次性讀出來。
count = socketChannel.read(bb);
if (count < 0) {
// 已經讀到流的結尾，或連接異常，需要關閉連接
socketChannel.close();
// 注意key.cancel()是在下次select()的時候纔會被清理
key.cancel();
return;
}
} catch (IOException e) {
e.printStackTrace();
}
// buffer的使用一定要好好看看API，buffer的熟練使用對NIO編程很重要
bb.flip();
int limit = bb.limit();
byte[] tmpbytes = new byte[limit];
bb.get(tmpbytes);
logger.debug("接受信息爲:" + new String(tmpbytes));
if (!isCache(key, tmpbytes)) {
byte[] bytes = (byte[]) key.attachment();
String requestStr = new String(bytes);
logger.debug("請求字符串:" + requestStr);
bb.clear();
if (requestStr.equals("gettime")) {
bb.put(new Date().toString().getBytes());
key.attach(new byte[0]);
} else if (requestStr.endsWith("clear")) {
key.attach(new byte[0]);
try {
bb.put("緩存已清理".getBytes("GB2312"));
} catch (UnsupportedEncodingException e) {
e.printStackTrace();
}
} else {
try {
bb.put("不能識別的請求".getBytes("GB2312"));
} catch (UnsupportedEncodingException e) {
e.printStackTrace();
}
}
bb.flip();
try {
socketChannel.write(bb);
} catch (IOException e) {
e.printStackTrace();
}
}

}

private boolean isCache(SelectionKey key, byte[] tmpbytes) {
Object obj = key.attachment();
byte[] bytes;
if (obj != null) {
bytes = (byte[]) obj;
} else {
bytes = new byte[0];
}
int sumLength = bytes.length + tmpbytes.length;
ByteBuffer bb = ByteBuffer.allocate(sumLength);
bb.put(bytes);
bb.put(tmpbytes);
bb.flip();
tmpbytes = bb.array();
if (tmpbytes[sumLength - 1] == 10) {
tmpbytes = new byte[sumLength - 2];
bb.get(tmpbytes);
key.attach(tmpbytes);
return false;
} else {
key.attach(tmpbytes);
return true;
}
}
}
}