使用原生NIO實現一個echo服務器

tcp的拆包處理使用的是定長解碼的方式。
服務器端:

public class EchoServer {
    public static final int port = 8888;

    public static void main(String[] args) throws IOException {
        Selector selector = Selector.open();

        ServerSocketChannel listener = ServerSocketChannel.open();
        // 綁定地址,並監聽
        listener.socket().bind(new InetSocketAddress("localhost", port));

        // 設置非阻塞
        listener.configureBlocking(false);
        // 註冊ACCPET事件
        listener.register(selector, SelectionKey.OP_ACCEPT);
        NIOServerConnection conn;
        while (true) {
            // 每1秒選擇一次
            if (selector.select(1000) == 0) {
                System.out.print(".");
                continue;
            }

            Iterator<SelectionKey> iter = selector.selectedKeys().iterator();

            while (iter.hasNext()) {
                SelectionKey key = iter.next();
                iter.remove();
                if (key.isAcceptable()) {
                    listener = (ServerSocketChannel) key.channel();
                    SocketChannel clientChannel = listener.accept();
                    // 設置地址複用
                    clientChannel.socket().setReuseAddress(true);
                    clientChannel.configureBlocking(false);
                    // 將接受的客戶端通道設置可讀
                    SelectionKey connKey = clientChannel.register(selector, SelectionKey.OP_READ);
                    // 使用該類處理讀寫請求
                    conn = new NIOServerConnection(connKey);
                    connKey.attach(conn);
                }

                if (key.isReadable()) {
                    conn = (NIOServerConnection) key.attachment();
                    conn.handleRead();
                }

                if (key.isValid() && key.isWritable()) {
                    conn = (NIOServerConnection) key.attachment();
                    conn.handleWrite();
                }

            }
        }

    }
}

服務器端處理讀寫類:

public class NIOServerConnection {
    private final SelectionKey key;
    private ByteBuffer data;
    private final ByteBuffer dataLengthBuffer;

    public NIOServerConnection(SelectionKey key) {
        this.key = key;
        dataLengthBuffer = ByteBuffer.allocate(4);
    }

    public void handleRead() {
        dataLengthBuffer.clear();
        SocketChannel channel = (SocketChannel) key.channel();
        try {
            long bytesRead = channel.read(dataLengthBuffer);
            assert bytesRead == 4;
            // 標記一下
            dataLengthBuffer.flip().mark();
            int len = dataLengthBuffer.getInt();
            // 在讀取的時候,還要重置的
            dataLengthBuffer.reset();
            data = ByteBuffer.allocate(len);
            bytesRead += channel.read(data);
            // 斷言全部讀取完畢
            assert bytesRead == (4 + len);
            byte[] b = data.array();
            int port = channel.socket().getPort();
            System.out.println("client port : " + port + " ; msg : " + new String(b));
            if (bytesRead == -1) {
                channel.close();
            } else {
                key.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
            }
        } catch (IOException e) {
            e.printStackTrace();
            try {
                channel.close();
            } catch (IOException e1) {
                e1.printStackTrace();
            }
        }
    }

    public void handleWrite() {
        data.flip();
        SocketChannel channel = (SocketChannel) key.channel();
        try {
            channel.write(new ByteBuffer[] { dataLengthBuffer, data });
        } catch (IOException e) {
            e.printStackTrace();
        }
        if (!data.hasRemaining())
            key.interestOps(SelectionKey.OP_READ);
        data.compact();
    }

}

客戶端:

/**
 * @author root TCP NIO 定長拆包
 */
public class EchoClient {
    public static final String[] messages = new String[] { "nio", "netty", "love", "java", "learn" };
    public static final Random rand = new Random();
    public static final String BLANK = " ";

    public static String getMsg() {
        StringBuilder msg = new StringBuilder();
        int len = messages.length;
        int words = rand.nextInt(len) + 1;
        for (int i = 0; i < words; i++) {
            msg.append(messages[rand.nextInt(len)] + BLANK);
        }
        return msg.toString() + getDate();
    }

    public static String getDate() {
        Date date = new Date();
        return date.toString();
    }

    public static void main(String[] args) {
        SocketChannel clientChannel = null;
        try {
            Selector selector = Selector.open();
            clientChannel = SocketChannel.open();
            clientChannel.connect(new InetSocketAddress("localhost", EchoServer.port));
            clientChannel.configureBlocking(false);
            SelectionKey key = clientChannel.register(selector, SelectionKey.OP_WRITE);
            NIOClientConnection conn = new NIOClientConnection(key);
            key.attach(conn);
            while (true) {
                Thread.sleep(2000);
                selector.select();
                Iterator<SelectionKey> iters = selector.selectedKeys().iterator();
                while (iters.hasNext()) {
                    key = iters.next();
                    iters.remove();
                    if (key.isWritable()) {
                        conn = (NIOClientConnection) key.attachment();
                        conn.doWrite(getMsg());
                    }
                    if (key.isReadable()) {
                        conn = (NIOClientConnection) key.attachment();
                        conn.doRead();
                    }
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
            try {
                clientChannel.close();
            } catch (IOException e1) {
                e1.printStackTrace();
            }
        }
    }

}

客戶端處理讀寫類:

public class NIOClientConnection {
    private final SelectionKey key;
    private ByteBuffer data;
    private final ByteBuffer dataLengthBuffer;

    public NIOClientConnection(SelectionKey key) {
        this.key = key;
        dataLengthBuffer = ByteBuffer.allocate(4);
    }

    public void doWrite(String msg) {
        dataLengthBuffer.clear();
        byte[] b = msg.getBytes();
        int len = b.length;
        // put操作後需要flip,wrap操作不用flip
        dataLengthBuffer.putInt(len);
        dataLengthBuffer.flip();
        data = ByteBuffer.wrap(b);
        SocketChannel channel = (SocketChannel) key.channel();
        try {
            System.out.println("send => " + msg);
            long bs = channel.write(new ByteBuffer[] { dataLengthBuffer, data });
            if (bs == 0)
                channel.close();
            else
                key.interestOps(SelectionKey.OP_READ);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public void doRead() {
        dataLengthBuffer.clear();
        data.clear();
        SocketChannel channel = (SocketChannel) key.channel();
        try {
            long bs = channel.read(new ByteBuffer[] { dataLengthBuffer, data });
            if (bs == 0 || bs == -1) {
                return;
            }
            dataLengthBuffer.flip();
            int len = dataLengthBuffer.getInt();
            byte[] buf = data.array();
            assert len + 4 == bs;
            System.out.println("receive => " + new String(buf));
            key.interestOps(SelectionKey.OP_WRITE);
        } catch (IOException e) {
            e.printStackTrace();
            try {
                channel.close();
            } catch (IOException e1) {
                e1.printStackTrace();
            }
        }
    }
}

測試:
服務器結果:
這裏寫圖片描述
客戶端:
這裏寫圖片描述

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章