java NIO解决沾包断包

网络不稳定、服务器分配的buffer不够存一条完整消息时服务端可能会出现一次read事件读出来的消息不是完整的，这个时候我们就需要自己粘包。

粘包思路：

客户端发送消息时先写入要发消息长度在后买面追加消息数据，服务接受消息时需要先解析出消息长度，然后分配一个消息长度大小的buffer，然后数据到buffer,读满buffer就正好读完一条消息了，这样无论消息长短都可以正确处理。

代码示列：

Server沾包代码：

if (key.isReadable()) {
            Message msg = (Message) key.attachment();
            if (msg == null) {
                msg = new Message();
                msg.head = ByteBuffer.allocate(4);
                key.attach(msg);
            }
            SocketChannel sc = (SocketChannel) key.channel();
            if (msg.head.hasRemaining()) {
                sc.read(msg.head);
                if (!msg.head.hasRemaining()) {
                    msg.bodyLen = bytesToInt(msg.head.array());
                    if (msg.body == null || msg.body.capacity() < msg.bodyLen) {
                        msg.body = ByteBuffer.allocate(msg.bodyLen);
                    } else {
                        msg.body.clear();
                        msg.body.limit(msg.bodyLen);
                    }
                }
            }
            if (msg.bodyLen > 0) {
                sc.read(msg.body);
                if (!msg.body.hasRemaining()) {
                    System.out.println("Server received" + sc.getRemoteAddress() + " data:" + new String(msg.body.array(), 0, msg.bodyLen));
                    doWrite(sc, "Server received");
                    msg.body.clear();
                    msg.head.clear();
                    msg.bodyLen = -1;
                }
            }
        }

Server完整代码：

import java.io.IOException;
import java.net.InetSocketAddress;
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.nio.charset.Charset;
import java.util.Iterator;

public class Server implements Runnable {
    private final int port;
    private volatile boolean stop;
    private ServerSocketChannel ssc;
    private Selector selector;


    public Server(int port) throws IOException {
        this.port = port;
        ssc = ServerSocketChannel.open();
        ssc.configureBlocking(false);
        ssc.socket().bind(new InetSocketAddress(port));
        selector = Selector.open();
        ssc.register(selector, SelectionKey.OP_ACCEPT);
    }

    @Override
    public void run() {
        while (!stop) {
            try {
                if (selector.selectNow() >= 0) {
                    Iterator iterator = selector.selectedKeys().iterator();
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        iterator.remove();
                        handleKey(key);
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    static class Message {
        ByteBuffer head;
        ByteBuffer body;
        int bodyLen = -1;
    }

    private void handleKey(SelectionKey key) throws IOException {
        if (key == null || !key.isValid()) {
            return;
        }

        if (key.isAcceptable()) {
            SocketChannel sc = ssc.accept();
            sc.configureBlocking(false);
            sc.register(selector, SelectionKey.OP_READ);
        }

        if (key.isReadable()) {
            Message msg = (Message) key.attachment();
            if (msg == null) {
                msg = new Message();
                msg.head = ByteBuffer.allocate(4);
                key.attach(msg);
            }
            SocketChannel sc = (SocketChannel) key.channel();
            if (msg.head.hasRemaining()) {
                sc.read(msg.head);
                if (!msg.head.hasRemaining()) {
                    msg.bodyLen = bytesToInt(msg.head.array());
                    if (msg.body == null || msg.body.capacity() < msg.bodyLen) {
                        msg.body = ByteBuffer.allocate(msg.bodyLen);
                    } else {
                        msg.body.clear();
                        msg.body.limit(msg.bodyLen);
                    }
                }
            }
            if (msg.bodyLen > 0) {
                sc.read(msg.body);
                if (!msg.body.hasRemaining()) {
                    System.out.println("Server received" + sc.getRemoteAddress() + " data:" + new String(msg.body.array(), 0, msg.bodyLen));
                    doWrite(sc, "Server received");
                    msg.body.clear();
                    msg.head.clear();
                    msg.bodyLen = -1;
                }
            }
        }
    }

    private void doWrite(SocketChannel sc, String s) throws IOException {
        sc.write(Charset.forName("utf-8").encode(s));
    }

    private int bytesToInt(byte[] head) {
        if (head == null || head.length != 4) {
            return 0;
        }
        return ((head[0] & 0xff) << 24) + ((head[1] & 0xff) << 16) + ((head[2] & 0xff) << 8) + ((head[3] & 0xff));
    }
}

client：

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.UUID;

public class Client implements Runnable {
    private static int cid;
    private int id;
    private final int port;
    private final String host;
    private boolean stop;
    private SocketChannel sc;
    private Selector selector;
    private ByteBuffer byteBuffer;

    public Client(String host, int port) throws IOException {
        id = ++cid;
        this.host = host;
        this.port = port;
        sc = SocketChannel.open();
        sc.configureBlocking(false);
        selector = Selector.open();
        byteBuffer = ByteBuffer.allocate(1024);
    }

    @Override
    public void run() {
        try {
            doConnect();
            while (!stop) {
                if (selector.select() > 0) {
                    Iterator iterator = selector.selectedKeys().iterator();
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        iterator.remove();
                        handleKey(key);
                    }
                }
            }
            selector.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void handleKey(SelectionKey key) throws IOException {
        if (key == null || !key.isValid()) {
            return;
        }

        if (key.isConnectable()) {
            if (sc.finishConnect()) {
                sc.register(selector, SelectionKey.OP_READ);
                String msg = "I am " + id + ", my data: " + id + ":" + UUID.randomUUID();
                if (Math.random() < 0.5) {
                    msg += UUID.randomUUID();
                }
                msg += ", end";
                doWrite(sc, msg);
                if (Math.random() > 0.5) {
                    msg = "I am " + id + ", Second message, end";
                    doWrite(sc, msg);
                }
            }
        }

        if (key.isReadable()) {
            byteBuffer.clear();
            int len = sc.read(byteBuffer);
            if (len > 0) {
                System.out.println(sc.getLocalAddress() + " client " + id + ", received:" + new String(byteBuffer.array(), 0, len));
                stop = true;
            }
        }
    }

    private void doWrite(SocketChannel sc, String s) throws IOException {
        ByteBuffer byteBuffer = ByteBuffer.allocate(4 + s.length());
        byte[] head = intToBytes(s.length());
        if (Math.random() < 0.5) {
            // Simulated network latency
            for (int i = 0; i < 4; i++) {
                sc.write(ByteBuffer.wrap(new byte[]{head[i]}));
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        } else {
            byteBuffer.put(head);
        }

        byteBuffer.put(s.getBytes("utf-8"));
        byteBuffer.flip();
        while (byteBuffer.hasRemaining()) {
            sc.write(ByteBuffer.wrap(new byte[]{byteBuffer.get()}));
            if (Math.random() < 0.4) {
                // Simulated network latency
                try {
                    Thread.sleep(100);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    private byte[] intToBytes(int length) {
        byte[] bytes = new byte[4];
        bytes[0] = (byte) (length >> 24 & 0xFF);
        bytes[1] = (byte) (length >> 16 & 0xFF);
        bytes[2] = (byte) (length >> 8 & 0xFF);
        bytes[3] = (byte) (length & 0xFF);
        return bytes;
    }

    private void doConnect() throws IOException {
        if (sc.isConnected() || sc.connect(new InetSocketAddress(host, port))) {
            sc.register(selector, SelectionKey.OP_READ);
        } else {
            sc.register(selector, SelectionKey.OP_CONNECT);
        }
    }
}

test：

import java.io.IOException;

public class Main {
    public static void main(String[] args) throws IOException, InterruptedException {
        Thread serverTheard = new Thread(new Server(8000));
        serverTheard.setDaemon(true);
        serverTheard.start();

        Thread.sleep(1000);
        for (int i = 0; i < 100; i++) {
            new Thread(new Client("localhost", 8000)).start();
        }
    }
}