Java Socket 短链接 自定义报文

最近对接的一个项目，使用 TCP Socket 短链接自定义报文格式，不是传统的 HTTP 接口方式对接，踩了一些坑，特此记录一下。特别是 TCP 是流式传输，会出现只收到前几个字节，后面的字节过一会才会收到的情况。

我这边的报文格式是，前4个字节是后面报文主体的长度，然后后面才是报文主体

SocketMessageReader.java 报文读取

import java.io.IOException;
import java.io.InputStream;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

// Socket 读报文，报文格式： 4字节报文长度+报文主体
public class SocketMessageReader implements AutoCloseable {

  private final Socket socket;
  private final InputStream in;

  public SocketMessageReader(Socket socket) throws IOException {
    this.socket = socket;
    this.in = socket.getInputStream();
  }

  /**
   * 读取一条完整报文
   */
  public String readMessage() throws IOException {
    // 1. 先读取4个字节，代表正文长度
    byte[] headerBytes = readFully(4);
    int contentLength = ByteBuffer.wrap(headerBytes).getInt(); // 默认大端序

    // 2. 再根据长度读取正文
    byte[] bodyBytes = readFully(contentLength);
    return new String(bodyBytes, StandardCharsets.UTF_8);
  }

  /**
   * 从输入流读取指定字节数，直到满为止（处理TCP分片问题）
   */
  private byte[] readFully(int length) throws IOException {
    byte[] buffer = new byte[length];
    int offset = 0;
    while (offset < length) {
      int read = in.read(buffer, offset, length - offset);
      if (read == -1) {
        throw new IOException("连接已关闭，未能完整读取数据");
      }
      offset += read;
    }
    return buffer;
  }

  @Override
  public void close() {
    try {
      in.close();
    } catch (Exception e) {
      e.printStackTrace();
    } finally {
      try {
        socket.close();
      } catch (Exception e) {
        e.printStackTrace();
      }
    }
  }
}

SocketMessageWriter.java 写入报文

import java.io.IOException;
import java.io.OutputStream;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;

// Socket 写报文，报文格式： 4字节报文长度+报文主体
public class SocketMessageWriter implements AutoCloseable {

  private final Socket socket;
  private final OutputStream out;

  private int length = -1;

  public SocketMessageWriter(Socket socket) throws IOException {
    this.socket = socket;
    this.out = socket.getOutputStream();
  }

  // 设置报文长度，只需调用一次，需搭配 writePartMsg 一起使用
  public void setLength(int length) throws IOException {
    this.length = length;
    byte[] headerBytes = ByteBuffer.allocate(4).putInt(length).array();
    out.write(headerBytes);
    out.flush();
  }

  // 发送部分消息，需提前调用  setLength 设置整体报文长度
  public void writePartMsg(String msg) throws IOException {
    if (length < 0) {
      throw new RuntimeException("请先设置length");
    }
    byte[] bodyBytes = msg.getBytes(StandardCharsets.UTF_8);
    out.write(bodyBytes);
    out.flush();
  }

  // 消息整体发送
  public void writeAllMsg(String body) throws IOException {
    byte[] bodyBytes = body.getBytes(StandardCharsets.UTF_8);
    byte[] headerBytes = ByteBuffer.allocate(4).putInt(bodyBytes.length).array();
    out.write(headerBytes);

    out.write(bodyBytes);
    out.flush();
  }


  @Override
  public void close() {
    try {
      out.close();
    } catch (Exception e) {
      e.printStackTrace();
    } finally {
      try {
        socket.close();
      } catch (IOException e) {
        e.printStackTrace();
      }
    }
  }

}

TcpServer 服务端

import java.io.IOException;
import java.net.ServerSocket;
import java.net.Socket;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class TcpServer {

  public static void main(String[] args) {
    ExecutorService executor = Executors.newCachedThreadPool();
    try (ServerSocket serverSocket = new ServerSocket(9000)) {
      System.out.println("服务端启动，等待连接...");
      while (true) {
        Socket socket = serverSocket.accept();
        executor.submit(() -> {
          String uuid = UUID.randomUUID().toString();
          try (SocketMessageReader reader = new SocketMessageReader(socket)) {
            while (true) {
              String message = reader.readMessage();
              System.out.println(uuid + " ==> 收到消息: " + message);
              try (SocketMessageWriter writer = new SocketMessageWriter(socket)) {
                writer.writeAllMsg(uuid + "服务端回复：" + message);
              }
            }
          } catch (IOException e) {
            System.out.println("客户端 " + uuid + " 已断开: " + e.getMessage());
          }
        });
      }
    } catch (Exception e) {
      e.printStackTrace();
    }
  }
}

TcpClient 客户端

import java.net.Socket;

public class TcpClient {

  public static void main(String[] args) {
    try (Socket socket = new Socket("127.0.0.1", 9000)) {
      SocketMessageWriter writer = new SocketMessageWriter(socket);

      writer.setLength(10);
      // 故意延时500ms，生产环境需要去掉
      Thread.sleep(500);

      writer.writePartMsg("hello");
      // 故意延时500ms，生产环境需要去掉
      Thread.sleep(500);

      writer.writePartMsg("world");

      SocketMessageReader socketMessageReader = new SocketMessageReader(socket);

      String s = socketMessageReader.readMessage();
      System.out.println("客户端接收：" + s);

    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }
}