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 都要高.
server | ws-tool | ws-tool async | uWebSocket | tungstenite |
---|---|---|---|---|
msg/sec | 5150± | 5150± | 3200± | 2850± |
测试机器为 i9-12900k + 32G 3200Mhz ddr4 内存.
ws-tool 项目功能矩阵
IO type | split | proxy | tls | deflate | use as client | use as server |
---|---|---|---|---|---|---|
blocking | ✅ | ✅ | ✅ | 🚧wip | ✅ | ✅ |
async | ✅ | ✅ | ✅ | 🚧wip | ✅ | ✅ |
项目来源与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, 方向或目标如下
-
- 换掉 tokio-codec, 自己维护协议解析状态, 这样的好处在于减少泛型且同步/异步接口可以共享代码(tokio-codec只有异步接口)
-
- 提供一套类似的, 涵盖最底层 frame 读写, 与最常用的 string, bytes 读写接口
-
- 允许用户在客户端发起请求, 处理服务器响应和服务器处理请求构造消息解析器时传入函数, 这样可以方便用户定义数据收发行为(如自动切包/合并小包, 限制最大载荷, 是否进行 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 之后发现两个主要的性能瓶颈, 分别是
- 有多次数据 copy, 特别是 load_test 的 payload 相对比较大, 多次 copy 带来不少性能损耗.
- 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 一下.