使用 NIO 搭建一个聊天室

使用 NIO 搭建一个聊天室

前面刚讲了使用 Socket 搭建了一个 Http Server，在最后我们使用了 NIO 对 Server 进行了优化，然后有小伙伴问到怎么使用 Socket 搭建聊天室，这节仍然使用 NIO 为基础进行搭建。

一、NIO 聊天室入门案例

该案例只有三个类：NioServer 聊天室服务端、NioClient 聊天室客户端、ClientThread 客户端线程。

服务端代码：

package com.fengsir.network.chatroom;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CharsetEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @Author FengZeng
 * @Date 2022-01-24 17:01
 * @Description Nio聊天室服务端
 */
public class NioServer {
  /**
   * 聊天室成员列表：
   */
  Map<String, SocketChannel> memberChannels;

  /**
   * 端口
   */
  private static final int PORT = 8000;

  /**
   * 选择器
   */
  private Selector selector;

  /**
   * 管道
   */
  private ServerSocketChannel server;

  /**
   * 缓冲
   */
  private ByteBuffer buffer;

  public NioServer() throws IOException {
    // 初始化 Selector 选择器
    this.selector = Selector.open();
    // 初始化 Channel 通道
    this.server = getServerChannel(selector);
    // 初始化 Buffer 缓冲：1k
    this.buffer = ByteBuffer.allocate(1024);
    // 初始化聊天室成员列表
    memberChannels = new ConcurrentHashMap<>();
  }

  /**
   * 初始化Channel通道
   *
   * @param selector 选择器
   * @return ServerSocketChannel
   * @throws IOException
   */
  private ServerSocketChannel getServerChannel(Selector selector) throws IOException {
    // 开辟一个 Channel 通道
    ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();

    // 通道设置为非阻塞模式
    serverSocketChannel.configureBlocking(false);

    // 通道注册绑定 Selector 选择器，通道中数据的事件类型为OP_ACCEPT
    serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

    // 通道绑定端口
    serverSocketChannel.socket().bind(new InetSocketAddress(PORT));

    return serverSocketChannel;
  }

  /**
   * 事件监听
   */
  public void listen() throws IOException {
    System.out.println("服务端启动......");
    try {
      // 据说用 while(true) 会多一个判断，用这种方式更好哈哈哈
      for (;;){
        // 作用：至少需要有一个事件发生，否则（如果count == 0）就继续阻塞循环
        int count = selector.select();
        if (count == 0) {
          continue;
        }
        // 获取 SelectorKey 的集合
        Set<SelectionKey> keySet = selector.selectedKeys();
        Iterator<SelectionKey> iterator = keySet.iterator();

        while (iterator.hasNext()) {
          // 当前事件对应的 SelectorKey
          SelectionKey key = iterator.next();

          // 删除当前事件：表示当前事件已经被消费了
          iterator.remove();

          // 接收事件已就绪：
          if (key.isAcceptable()) {

            // 通过key获取ServerSocketChannel
            ServerSocketChannel server = (ServerSocketChannel) key.channel();

            // 通过 ServerSocketChannel 获取SocketChannel
            SocketChannel channel = server.accept();

            // channel 设置为非阻塞模式
            channel.configureBlocking(false);
            // channel 绑定选择器，当前事件切换为 读就绪
            channel.register(selector, SelectionKey.OP_READ);

            // 从channel中获取Host、端口等信息
            System.out.println("客户端连接："
                + channel.socket().getInetAddress().getHostName() + ":"
                + channel.socket().getPort());

            // Read就绪事件
          } else if (key.isReadable()) {

            SocketChannel channel = (SocketChannel) key.channel();
            // 用于解密消息内容
            CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder();

            // 将消息数据从通道 channel 读取到缓冲buffer
            buffer.clear();
            channel.read(buffer);
            buffer.flip();
            // 获取解密后的消息内容：
            String msg = decoder.decode(buffer).toString();
            if (!"".equals(msg)) {
              System.out.println("收到：" + msg);
              if (msg.startsWith("username:")) {
                String username = msg.replaceAll("username:", "");
                memberChannels.put(username, channel);
                System.out.println("用户总数：" + memberChannels.size());
              } else {
                // 转发消息给客户端
                String[] arr = msg.split(":");
                if (arr.length == 3) {
                  // 发送者
                  String from = arr[0];
                  // 接收者
                  String to = arr[1];
                  // 发送内容
                  String content = arr[2];
                  System.out.println(from + " 发送给 " + to + " 的消息：" + content);

                  if (memberChannels.containsKey(to)) {
                    // 解密
                    CharsetEncoder encoder = StandardCharsets.UTF_8.newEncoder();
                    // 给接收者发送消息
                    memberChannels.get(to).write(encoder.encode(CharBuffer.wrap(from + ":" + content)));
                  }
                }
              }
            }

          }
        }
      }
    }catch (Exception e){
      System.out.println("服务端启动失败......");
      e.printStackTrace();
    }finally {
      try {
        // 先关闭选择器，在关闭通道
        // 调用 close() 方法将会关闭Selector，同时也会将关联的SelectionKey失效，但不会关闭Channel。
        selector.close();
        server.close();
      } catch (IOException e) {
        e.printStackTrace();
      }
    }
  }

  public static void main(String[] args) throws IOException {
    // 服务端启动：
    new NioServer().listen();
  }

}

客户端线程类：

package com.fengsir.network.chatroom;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CharsetEncoder;
import java.nio.charset.StandardCharsets;
import java.util.Iterator;

/**
 * @Author FengZeng
 * @Date 2022-01-24 17:16
 * @Description Nio聊天室客户端线程
 */
public class ClientThread extends Thread {
  /**
   * 解密
   */
  private CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder();

  /**
   * 加密
   */
  private CharsetEncoder encoder = StandardCharsets.UTF_8.newEncoder();

  /**
   * 选择器
   */
  private Selector selector = null;

  /**
   * 通道
   */
  private SocketChannel socket = null;

  /**
   * 通道key
   */
  private SelectionKey clientKey = null;

  /**
   * 用户名
   */
  private String username;

  public ClientThread(String username) {
    try {
      // 创建一个Selector
      selector = Selector.open();

      // 创建Socket并注册
      socket = SocketChannel.open();
      socket.configureBlocking(false);
      clientKey = socket.register(selector, SelectionKey.OP_CONNECT);

      // 连接到远程地址
      InetSocketAddress ip = new InetSocketAddress("localhost", 8000);
      socket.connect(ip);

      this.username = username;
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  /**
   * 开辟读取事件的线程
   */
  @Override
  public void run() {
    try {
      // 监听事件（无限循环）
      for (; ; ) {
        // 监听事件
        int count = selector.select();
        if (count == 0) {
          continue;
        }
        // 事件来源列表
        Iterator<SelectionKey> it = selector.selectedKeys().iterator();

        while (it.hasNext()) {
          SelectionKey key = it.next();
          // 删除当前事件
          it.remove();

          // 判断事件类型
          if (key.isConnectable()) {
            // 连接事件
            SocketChannel channel = (SocketChannel) key.channel();

            if (channel.isConnectionPending()) {
              channel.finishConnect();
            }
            channel.register(selector, SelectionKey.OP_READ);
            System.out.println("连接服务器端成功！");

            // 发送用户名
            send("username:" + this.username);
          } else if (key.isReadable()) {
            // 读取数据事件
            SocketChannel channel = (SocketChannel) key.channel();

            // 读取数据
            ByteBuffer buffer = ByteBuffer.allocate(50);
            channel.read(buffer);
            buffer.flip();
            String msg = decoder.decode(buffer).toString();
            System.out.println("收到: " + msg);
          }
        }
      }
    } catch (IOException e) {
      e.printStackTrace();
    } finally {
      // 关闭
      try {
        selector.close();
        socket.close();
      } catch (IOException ignored) {
      }
    }
  }

  /**
   * 发送消息
   *
   * @param msg message
   */
  public void send(String msg) {
    try {
      SocketChannel client = (SocketChannel) clientKey.channel();
      client.write(encoder.encode(CharBuffer.wrap(msg)));
    } catch (Exception e) {
      e.printStackTrace();
    }
  }

  /**
   * 关闭客户端
   */
  public void close() {
    try {
      selector.close();
      socket.close();
    } catch (IOException ignored) {
    }
  }

}

客户端类：

package com.fengsir.network.chatroom;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

/**
 * @Author FengZeng
 * @Date 2022-01-24 17:20
 * @Description TODO
 */
public class NioClient {
  public static void main(String[] args) {
    // 当前客户端的用户名
    String username = "fengzeng";
    // 为当前客户端开辟一个线程
    ClientThread client = new ClientThread(username);
    client.start();

    // 输入输出流
    BufferedReader bfReader = new BufferedReader(new InputStreamReader(System.in));

    try {
      // 循环读取键盘输入
      String readline;
      while ((readline = bfReader.readLine()) != null) {
        if ("bye".equals(readline)) {
          client.close();
          System.exit(0);
        }
        // 发送消息
        client.send(username + ":" + readline);
      }
    } catch (IOException e) {
      e.printStackTrace();
    }
  }
}

运行效果图就是这样子，代码结合注释看，应该都能理解。