ws-tool 0.5.0 - 项目设计与优化小结

今天在完成 ws-tool 代理支持的的 PR 后 ws-tool 0.5.0 正式发布.

项目简介

ws-tool 一个 rust 实现的 websocket 库, 你可以用它实现 websocket 客户端或服务端, ws-tool 支持同步和异步接口, 支持将一个读写连接分割成 read half 和 write half, 可以充分使用 tcp 的全双工模式通信, 而且 ws-tool 也拥有不错的性能.

使用 cpp uWebSocket 提供的 load_test 工具测试单线程下 msg/sec 指标, ws-tool 的同步异步接口性能比 uWebSocket 和 tungstenite 都要高.

测试机器为 i9-12900k + 32G 3200Mhz ddr4 内存.

ws-tool 项目功能矩阵

项目来源与API设计

ws-tool 项目来源于之前在使用 rust 实现币安 websocket 接口行情数据收集时发现 tungstenite 不支持代理, 在我 fork 项目并实现功能后, tungstenite 没有接受我的 PR. 因此如果我继续使用 tungstenite, 我必须维护自己的分支, 那为什么不自己写一个呢? 而且 tungstenite 的异步绑定接口不太好用, 我希望有一个支持同步/异步接口, http/socks5 代理的 websocket 库, ws-tool 由此诞生.

从最初 0.2.0 到 0.5.0 版本, ws-tool 经历了2次重构. 初版 ws-tool 只能保证 websocket 协议实现正确, 使用 tokio/codec 实现协议解析, 导致泛型太多, API 难用.

因此第一次重构主要方向是调整 API, 方向或目标如下

• 1. 换掉 tokio-codec, 自己维护协议解析状态, 这样的好处在于减少泛型且同步/异步接口可以共享代码(tokio-codec只有异步接口)
• 2. 提供一套类似的, 涵盖最底层 frame 读写, 与最常用的 string, bytes 读写接口
• 3. 允许用户在客户端发起请求, 处理服务器响应和服务器处理请求构造消息解析器时传入函数, 这样可以方便用户定义数据收发行为(如自动切包/合并小包, 限制最大载荷, 是否进行 utf-8 校验和自定义扩展等)

在 ws-tool 在 0.3.0 基本上实现了上述重构目标.

首先 ws-tool 提供了从类似的同步/异步构造和数据读写接口, 以项目的 example 为例.

同步 server 构造和读写消息

let mut server = ServerBuilder::accept(
    stream,
    default_handshake_handler,
    WsStringCodec::factory,
)
.unwrap();

loop {
    match server.receive() {
        Ok(msg) => server.send((msg.code, msg.data)).unwrap(),
        Err(e) => {
            dbg!(e);
            break;
        }
    }
}

异步 serer 构造读写消息

let mut server = ServerBuilder::async_accept(
    stream,
    default_handshake_handler,
    AsyncWsStringCodec::factory,
)
.await
.unwrap();

loop {
    match server.receive().await {
        Ok(msg) => {
            server.send((msg.code, msg.data)).await.unwrap()
        }
        Err(e) => {
            dbg!(e);
            break;
        }
    }
}

可以看到两者API除了异步接口时需要 async_accept 和 await future, 在使用上基本没有区别.

对于数据写入接口的设计可以查看Into特性实现参数多态的技巧这篇文章.

在 builder accept 时可以看到除了 stream 还需要传入两个参数, 这两个参数是实现重构目标3的关键. 对于 server 来说, 它可能会检查 websocket 请求 header 之外额外的 header, 并决定 server 是否接受 upgrade 请求.

如果查看第二个参数的类型, 可以发现这是一个 Fn(http::Request<()>) -> Result<(http::Request<()>, http::Response<T>), WsError>, 即参数为 client 握手请求, 并返回该请求(后续构建codec时还需要)和response. 用 户只要传入符合签名的函数即实现自定义握手请求处理逻辑. 当然为了使用简便, ws-tool 提供了 default_handshake_handler 函数以应对大多数情况.

第三个参数类型为 Fn(http::Request<()>, WsStream<S>) -> Result<C, WsError>, 这个函数接受 client 握手请求(可能包含认证信息等其他有用信息)以及已经把第二个参数返回返回个client的 stream, 即已经完成握手的 stream, 你可以在上面读写 websocket frame. 用户可以在 stream 上包装自己逻辑, 实现自定义的 codec, 如根据 header 解密 bytes 信息等.

为了方便使用, ws-tool 提供了 FrameCodec, StringCodec, BytesCodec 三种最基本也是最常用的 codec, 通过它们, 用户可以方便的读写最原始的 frame, 或者更高级的 text/bytes.

性能优化

有了相对好用的 API, 下一步是性能. 在最初的测试里, ws-tool 是所有比较对象中性能最低的. 使用 cargo-flamegraph profile 之后发现两个主要的性能瓶颈, 分别是

1. 有多次数据 copy, 特别是 load_test 的 payload 相对比较大, 多次 copy 带来不少性能损耗.
2. mask 数据实现性能低下, 在大 payload 下严重拖累性能.

有了目标后, 就可以着手进行优化.

为了减少数据 copy, ws-tool 在读写时不再先写入固定大小的 buffer, 而是利用 BytesMut 的 API, 操作 len, 记录偏移量等方式将数据直接写入 frame 底层 buffer 中, 具体实现可以参考项目代码. 但这样做在写入 frame 时又会带来一个问题, 这种 frame 需要“拥有“数据(BytesMut独占所有权), 但对于 server 端, 因为其不需要 mask payload 的特性, 提供一个 &[u8] 甚至 Vec<&[u8]> 类型的 payload 显然限制更少, 用着更方便, 而且避免了把数据 copy 到 BytesMut 的性能损耗.

所以 ws-tool 又对底层 Frame 数据类型进行了一次重构, 目前共分出三种 frame

• ReadFrame, 底层是一个 BytesMut, 来自于从 stream 读到的 frame, frame 多大就多大, 没有占用额外空间存储信息.
• BorrowedFrame, 有 FrameHeader 和 Vec<&[u8]> 组成, 没有 payload 所有权, 多数情况用于写入
• OwnedFrame, 底层是 header 和 payload, 拥有 payload 所有权(都是 BytesMut), 这个类型是为了给用户操作底层 frame 数据提供方便的接口

至于 mask 可以直接将 tungstenite 实现 copy 过来. 以后可以考虑使用 std::simd 来提供 portable simd 加速, 理论上能提供更好的 mask 性能.

最后的测试结果表明, 这些优化非常有效, 现在 ws-tool 性能是最好的.

回归初心

ws-tool 最初目标之一是支持 proxy, 可惜 crates.io 中 http/socks5 代理绝大部分只支持 async, 导致 ws-tool 在同步接口中无法使用 proxy. 为此我又实现了 http/socks5 代理客户端, 并发布在 crates.io 上.

自此, ws-tool 已经实现 sync/async 接口除 deflate 扩展外的设计目标.

总结

感觉没什么好总结的, 设计和实现都在代码里, 如果觉得文章或项目有用请点个赞或 star 一下.