Skip to content

GitLab

Unverified Commit 63d7c01c authored by Ryan Olson's avatar Ryan Olson Committed by GitHub
Browse files

feat: velo-backend (#6547) 


Signed-off-by: default avatarRyan Olson <rolson@nvidia.com>
parent 2cc92bfa
Show whitespace changes
Inline Side-by-side
Showing with 3538 additions and 45 deletions
Cargo.lock
View file @ 63d7c01c
......@@ -306,6 +306,25 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "async-ucx"
version = "0.2.0"
source = "git+https://github.com/ryanolson/async-ucx.git?branch=ryan%2Fucx#fbcb8e3d281d7c910e3dcd84f7eb3e8ee253ac6e"
dependencies = [
"bytes",
"crossbeam",
"derivative",
"futures",
"futures-lite",
"lazy_static",
"log",
"pin-project",
"socket2 0.4.10",
"thiserror 1.0.69",
"tokio",
"ucx1-sys",
]
[[package]]
name = "async_zmq"
version = "0.4.0"
......@@ -584,7 +603,7 @@ version = "1.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
dependencies = [
"fastrand",
"fastrand 2.3.0",
"gloo-timers",
"tokio",
]
......@@ -642,6 +661,29 @@ dependencies = [
"virtue",
]
[[package]]
name = "bindgen"
version = "0.66.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f2b84e06fc203107bfbad243f4aba2af864eb7db3b1cf46ea0a023b0b433d2a7"
dependencies = [
"bitflags 2.11.0",
"cexpr",
"clang-sys",
"lazy_static",
"lazycell",
"log",
"peeking_take_while",
"prettyplease",
"proc-macro2",
"quote",
"regex",
"rustc-hash 1.1.0",
"shlex",
"syn 2.0.117",
"which",
]
[[package]]
name = "bindgen"
version = "0.70.1"
......@@ -1621,6 +1663,17 @@ dependencies = [
"serde_core",
]
[[package]]
name = "derivative"
version = "2.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcc3dd5e9e9c0b295d6e1e4d811fb6f157d5ffd784b8d202fc62eac8035a770b"
dependencies = [
"proc-macro2",
"quote",
"syn 1.0.109",
]
[[package]]
name = "derive-getters"
version = "0.5.0"
......@@ -2385,6 +2438,15 @@ dependencies = [
"regex-syntax",
]
[[package]]
name = "fastrand"
version = "1.9.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e51093e27b0797c359783294ca4f0a911c270184cb10f85783b118614a1501be"
dependencies = [
"instant",
]
[[package]]
name = "fastrand"
version = "2.3.0"
......@@ -2511,7 +2573,7 @@ version = "0.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be"
dependencies = [
"fastrand",
"fastrand 2.3.0",
"futures-core",
"futures-sink",
"spin",
......@@ -2643,6 +2705,21 @@ version = "0.3.32"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
[[package]]
name = "futures-lite"
version = "1.13.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49a9d51ce47660b1e808d3c990b4709f2f415d928835a17dfd16991515c46bce"
dependencies = [
"fastrand 1.9.0",
"futures-core",
"futures-io",
"memchr",
"parking",
"pin-project-lite",
"waker-fn",
]
[[package]]
name = "futures-macro"
version = "0.3.32"
......@@ -3580,9 +3657,9 @@ checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
[[package]]
name = "jiff"
version = "0.2.22"
version = "0.2.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "819b44bc7c87d9117eb522f14d46e918add69ff12713c475946b0a29363ed1c2"
checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359"
dependencies = [
"jiff-static",
"jiff-tzdb-platform",
......@@ -3595,9 +3672,9 @@ dependencies = [
[[package]]
name = "jiff-static"
version = "0.2.22"
version = "0.2.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "470252db18ecc35fd766c0891b1e3ec6cbbcd62507e85276c01bf75d8e94d4a1"
checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4"
dependencies = [
"proc-macro2",
"quote",
......@@ -3606,9 +3683,9 @@ dependencies = [
[[package]]
name = "jiff-tzdb"
version = "0.1.5"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68971ebff725b9e2ca27a601c5eb38a4c5d64422c4cbab0c535f248087eda5c2"
checksum = "c900ef84826f1338a557697dc8fc601df9ca9af4ac137c7fb61d4c6f2dfd3076"
[[package]]
name = "jiff-tzdb-platform"
......@@ -3969,6 +4046,12 @@ version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
[[package]]
name = "lazycell"
version = "1.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55"
[[package]]
name = "leb128fmt"
version = "0.1.0"
......@@ -4311,9 +4394,9 @@ dependencies = [
[[package]]
name = "minijinja"
version = "2.16.0"
version = "2.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c54f3bcc034dd74496b5ca929fd0b710186672d5ff0b0f255a9ceb259042ece"
checksum = "65ab6f50e4e8fb40bd21f527066bd019f5b029035b4e5ac9b9f9ba526c6bd87b"
dependencies = [
"memo-map",
"self_cell",
......@@ -4322,9 +4405,9 @@ dependencies = [
[[package]]
name = "minijinja-contrib"
version = "2.16.0"
version = "2.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "86201a560fde90bf270d3c9d5da0c2a5a0c25bd3242d674fcf03e6a35fcfa7ba"
checksum = "2db46838f53a8c6e1d19135432e6b9bdbf4f0382c4fbf163f7f9dc24e31800e1"
dependencies = [
"minijinja",
"serde",
......@@ -5343,6 +5426,12 @@ dependencies = [
"syn 2.0.117",
]
[[package]]
name = "peeking_take_while"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099"
[[package]]
name = "pem"
version = "3.0.6"
......@@ -5972,9 +6061,9 @@ dependencies = [
[[package]]
name = "quote"
version = "1.0.44"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
......@@ -6941,6 +7030,16 @@ dependencies = [
"serde",
]
[[package]]
name = "serde_bytes"
version = "0.11.19"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8"
dependencies = [
"serde",
"serde_core",
]
[[package]]
name = "serde_cbor"
version = "0.11.2"
......@@ -7241,6 +7340,16 @@ version = "1.15.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
[[package]]
name = "socket2"
version = "0.4.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7916fc008ca5542385b89a3d3ce689953c143e9304a9bf8beec1de48994c0d"
dependencies = [
"libc",
"winapi 0.3.9",
]
[[package]]
name = "socket2"
version = "0.5.10"
......@@ -7483,7 +7592,7 @@ version = "3.26.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "82a72c767771b47409d2345987fda8628641887d5466101319899796367354a0"
dependencies = [
"fastrand",
"fastrand 2.3.0",
"getrandom 0.4.2",
"once_cell",
"rustix 1.1.4",
......@@ -8281,6 +8390,15 @@ version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971"
[[package]]
name = "ucx1-sys"
version = "0.2.0"
source = "git+https://github.com/ryanolson/async-ucx.git?branch=ryan%2Fucx#fbcb8e3d281d7c910e3dcd84f7eb3e8ee253ac6e"
dependencies = [
"bindgen 0.66.1",
"pkg-config",
]
[[package]]
name = "unarray"
version = "0.1.4"
......@@ -8616,6 +8734,20 @@ version = "0.2.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426"
[[package]]
name = "velo-common"
version = "1.0.0"
dependencies = [
"bytes",
"rmp-serde",
"serde",
"serde_bytes",
"serde_json",
"thiserror 2.0.18",
"uuid",
"xxhash-rust",
]
[[package]]
name = "velo-events"
version = "1.0.0"
......@@ -8632,6 +8764,43 @@ dependencies = [
"xxhash-rust",
]
[[package]]
name = "velo-transports"
version = "1.0.0"
dependencies = [
"anyhow",
"async-nats",
"async-ucx",
"axum 0.8.4",
"bs58",
"bytes",
"dashmap 6.1.0",
"flume",
"futures",
"http 1.4.0",
"http-body 1.0.1",
"http-body-util",
"hyper 1.8.1",
"hyper-util",
"nix 0.30.1",
"parking_lot",
"prost 0.13.5",
"reqwest 0.12.28",
"rmp-serde",
"serde_json",
"socket2 0.6.2",
"thiserror 2.0.18",
"tokio",
"tokio-stream",
"tokio-util",
"tonic 0.13.1",
"tonic-build 0.13.1",
"tower 0.5.3",
"tracing",
"tracing-subscriber",
"velo-common",
]
[[package]]
name = "version-compare"
version = "0.2.1"
......@@ -8670,6 +8839,12 @@ dependencies = [
"libc",
]
[[package]]
name = "waker-fn"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "317211a0dc0ceedd78fb2ca9a44aed3d7b9b26f81870d485c07122b4350673b7"
[[package]]
name = "walkdir"
version = "2.5.0"
......@@ -8863,6 +9038,18 @@ version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a28ac98ddc8b9274cb41bb4d9d4d5c425b6020c50c46f25559911905610b4a88"
[[package]]
name = "which"
version = "4.4.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "87ba24419a2078cd2b0f2ede2691b6c66d8e47836da3b6db8265ebad47afbfc7"
dependencies = [
"either",
"home",
"once_cell",
"rustix 0.38.44",
]
[[package]]
name = "winapi"
version = "0.2.8"
......
Cargo.toml
View file @ 63d7c01c
......@@ -19,7 +19,8 @@ members = [
"lib/bench",
"lib/bindings/c",
"lib/bindings/python/codegen",
"lib/config",
"lib/velo-common",
"lib/velo-transports",
"lib/velo-events",
]
resolver = "3"
......@@ -53,6 +54,8 @@ kvbm-logical = { path = "lib/kvbm-logical", version = "1.0.0" }
kvbm-physical = { path = "lib/kvbm-physical", version = "1.0.0" }
# velo
velo-common = { path = "lib/velo-common", version = "1.0.0" }
velo-transports = { path = "lib/velo-transports", version = "1.0.0" }
velo-events = { path = "lib/velo-events", version = "1.0.0" }
# External dependencies
......@@ -99,6 +102,7 @@ reqwest = { version = "0.12.24", default-features = false, features = [
"rustls-tls",
] }
rmp-serde = { version = "1" }
serde_bytes = { version = "0.11" }
# "rc" is for async-openai. Allows serializing Rc and Arc. Generally avoid doing that.
serde = { version = "1", features = ["derive", "rc"] }
serde_json = { version = "1" }
......
dynamo.code-workspace
View file @ 63d7c01c
......@@ -2,9 +2,6 @@
"folders": [
{
"path": "."
},
{
"path": "../dynamo-tpm"
}
],
"settings": {
......
lib/bindings/kvbm/Cargo.lock
View file @ 63d7c01c
......@@ -2965,9 +2965,9 @@ checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
[[package]]
name = "jiff"
version = "0.2.22"
version = "0.2.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "819b44bc7c87d9117eb522f14d46e918add69ff12713c475946b0a29363ed1c2"
checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359"
dependencies = [
"jiff-static",
"jiff-tzdb-platform",
......@@ -2980,9 +2980,9 @@ dependencies = [
[[package]]
name = "jiff-static"
version = "0.2.22"
version = "0.2.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "470252db18ecc35fd766c0891b1e3ec6cbbcd62507e85276c01bf75d8e94d4a1"
checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4"
dependencies = [
"proc-macro2",
"quote",
......@@ -2991,9 +2991,9 @@ dependencies = [
[[package]]
name = "jiff-tzdb"
version = "0.1.5"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68971ebff725b9e2ca27a601c5eb38a4c5d64422c4cbab0c535f248087eda5c2"
checksum = "c900ef84826f1338a557697dc8fc601df9ca9af4ac137c7fb61d4c6f2dfd3076"
[[package]]
name = "jiff-tzdb-platform"
......@@ -3584,9 +3584,9 @@ dependencies = [
[[package]]
name = "minijinja"
version = "2.16.0"
version = "2.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c54f3bcc034dd74496b5ca929fd0b710186672d5ff0b0f255a9ceb259042ece"
checksum = "65ab6f50e4e8fb40bd21f527066bd019f5b029035b4e5ac9b9f9ba526c6bd87b"
dependencies = [
"memo-map",
"self_cell",
......@@ -3595,9 +3595,9 @@ dependencies = [
[[package]]
name = "minijinja-contrib"
version = "2.16.0"
version = "2.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "86201a560fde90bf270d3c9d5da0c2a5a0c25bd3242d674fcf03e6a35fcfa7ba"
checksum = "2db46838f53a8c6e1d19135432e6b9bdbf4f0382c4fbf163f7f9dc24e31800e1"
dependencies = [
"minijinja",
"serde",
......@@ -5186,9 +5186,9 @@ dependencies = [
[[package]]
name = "quote"
version = "1.0.44"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
......
lib/bindings/python/Cargo.lock
View file @ 63d7c01c
......@@ -3026,9 +3026,9 @@ checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2"
[[package]]
name = "jiff"
version = "0.2.22"
version = "0.2.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "819b44bc7c87d9117eb522f14d46e918add69ff12713c475946b0a29363ed1c2"
checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359"
dependencies = [
"jiff-static",
"jiff-tzdb-platform",
......@@ -3041,9 +3041,9 @@ dependencies = [
[[package]]
name = "jiff-static"
version = "0.2.22"
version = "0.2.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "470252db18ecc35fd766c0891b1e3ec6cbbcd62507e85276c01bf75d8e94d4a1"
checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4"
dependencies = [
"proc-macro2",
"quote",
......@@ -3052,9 +3052,9 @@ dependencies = [
[[package]]
name = "jiff-tzdb"
version = "0.1.5"
version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68971ebff725b9e2ca27a601c5eb38a4c5d64422c4cbab0c535f248087eda5c2"
checksum = "c900ef84826f1338a557697dc8fc601df9ca9af4ac137c7fb61d4c6f2dfd3076"
[[package]]
name = "jiff-tzdb-platform"
......@@ -3633,9 +3633,9 @@ dependencies = [
[[package]]
name = "minijinja"
version = "2.16.0"
version = "2.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5c54f3bcc034dd74496b5ca929fd0b710186672d5ff0b0f255a9ceb259042ece"
checksum = "65ab6f50e4e8fb40bd21f527066bd019f5b029035b4e5ac9b9f9ba526c6bd87b"
dependencies = [
"memo-map",
"self_cell",
......@@ -3644,9 +3644,9 @@ dependencies = [
[[package]]
name = "minijinja-contrib"
version = "2.16.0"
version = "2.17.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "86201a560fde90bf270d3c9d5da0c2a5a0c25bd3242d674fcf03e6a35fcfa7ba"
checksum = "2db46838f53a8c6e1d19135432e6b9bdbf4f0382c4fbf163f7f9dc24e31800e1"
dependencies = [
"minijinja",
"serde",
......@@ -5245,9 +5245,9 @@ dependencies = [
[[package]]
name = "quote"
version = "1.0.44"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
......
lib/runtime/examples/Cargo.lock
View file @ 63d7c01c
......@@ -2713,9 +2713,9 @@ dependencies = [
[[package]]
name = "quote"
version = "1.0.44"
version = "1.0.45"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4"
checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924"
dependencies = [
"proc-macro2",
]
......
lib/velo-common/Cargo.toml 0 → 100644
View file @ 63d7c01c
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
[package]
name = "velo-common"
version.workspace = true
edition.workspace = true
description = "Common types for Velo distributed systems (identity, address, transport)"
authors.workspace = true
license.workspace = true
homepage.workspace = true
repository.workspace = true
keywords.workspace = true
[dependencies]
bytes = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
thiserror = { workspace = true }
uuid = { workspace = true, features = ["serde", "v4"] }
xxhash-rust = { workspace = true }
rmp-serde = { workspace = true }
serde_bytes = { workspace = true }
lib/velo-common/README.md 0 → 100644
View file @ 63d7c01c
# velo-common
Common types for the Velo distributed systems stack.
## Overview
This crate provides the foundational types used across Velo for identity and addressing. The design prioritizes:
- **Compact representations** for embedding in fixed-size handles
- **Transport-agnostic addressing** without enumerating all possible transports
- **KV-store friendly** serialization using opaque `Bytes`
## Identity Types
### InstanceId
`InstanceId` is a UUID-based identifier that serves as the **source of truth** for identifying a running Velo instance. It is used for:
- Transport-level routing
- Discovery registration
- Peer management
```rust
use velo_common::InstanceId;
let instance_id = InstanceId::new_v4();
let uuid = instance_id.as_uuid();
```
### WorkerId
`WorkerId` is a deterministic 64-bit identifier derived from `InstanceId` via xxh3 hash. The compact representation enables embedding worker identity into fixed-size handles.
**Design rationale**: A `u128` handle can encode:
- 64 bits for `WorkerId`
- 64 bits for additional data (sequence numbers, flags, etc.)
This value-semantics approach simplifies passing identity through systems that work with fixed-size integers.
```rust
use velo_common::InstanceId;
let instance_id = InstanceId::new_v4();
let worker_id = instance_id.worker_id(); // Deterministic derivation
// Embed in a u128 handle
let other_data: u128 = 42;
let handle: u128 = (worker_id.as_u64() as u128) << 64 | other_data;
```
The derivation is always consistent—calling `worker_id()` multiple times returns the same value.
## Address Types
### WorkerAddress
`WorkerAddress` is an opaque byte container holding transport endpoint information. Internally, it's a MessagePack-encoded map of `TransportKey -> Bytes`, but this structure is intentionally hidden from consumers.
**Key design decisions**:
1. **Opaque values**: Transport endpoints are stored as raw bytes. They could be simple strings (`"tcp://127.0.0.1:5555"`) or complex serialized objects. The interpretation is left to the transport implementation.
2. **No transport enum**: Rather than defining an enum of all possible transports with their configurations, we use string keys (`"tcp"`, `"rdma"`, `"grpc"`, etc.). This allows transports to be added without modifying the common types.
3. **KV-store friendly**: The entire address serializes to a `Bytes` blob, suitable for storage in etcd, Redis, or any key-value store without schema changes.
```rust
use velo_common::WorkerAddress;
use std::collections::HashMap;
// Addresses are typically constructed by velo-transports transport builders,
// but here we demonstrate the consumer API:
let mut map = HashMap::new();
map.insert("tcp".to_string(), b"tcp://127.0.0.1:5555".to_vec());
map.insert("rdma".to_string(), b"rdma://10.0.0.1:6666".to_vec());
let encoded = rmp_serde::to_vec(&map).unwrap();
let address = WorkerAddress::from_encoded(encoded);
// Reading an address (consumer perspective)
let transports = address.available_transports().unwrap(); // ["tcp", "rdma"]
let tcp_endpoint = address.get_entry("tcp").unwrap(); // Some(Bytes)
```
### TransportKey
A type-safe wrapper around transport identifiers. Provides zero-cost abstraction over `Arc<str>` with efficient cloning and HashMap compatibility.
```rust
use velo_common::TransportKey;
use std::collections::HashMap;
let key = TransportKey::from("tcp");
let key2: TransportKey = "rdma".into();
// Works with HashMap lookups via Borrow<str>
let mut map = HashMap::new();
map.insert(TransportKey::from("tcp"), "127.0.0.1:5555");
assert!(map.get("tcp").is_some()); // &str lookup works
```
### PeerInfo
Combines `InstanceId` and `WorkerAddress` into a single structure representing a discoverable peer. This is the primary type exchanged during peer discovery and registration.
```rust
use velo_common::{InstanceId, PeerInfo, WorkerAddress};
use std::collections::HashMap;
let instance_id = InstanceId::new_v4();
let map: HashMap<String, Vec<u8>> = HashMap::new();
let encoded = rmp_serde::to_vec(&map).unwrap();
let worker_address = WorkerAddress::from_encoded(encoded);
let peer_info = PeerInfo::new(instance_id, worker_address);
assert_eq!(peer_info.instance_id(), instance_id);
```
## Address Construction
`WorkerAddress` instances are constructed by transport builders in `velo-transports`. Each transport (TCP, gRPC, NATS, UCX, etc.) contributes its endpoint data:
```text
┌─────────────────┐ ┌─────────────────┐ ┌─────────────────┐
│ TCP Transport │ │ gRPC Transport │ │ UCX Transport │
│ Builder │ │ Builder │ │ Builder │
└────────┬────────┘ └────────┬────────┘ └────────┬────────┘
│ │ │
│ "tcp" -> endpoint │ "grpc" -> endpoint │ "ucx" -> blob
│ │ │
└───────────────────────┼───────────────────────┘
│
▼
┌────────────────────────┐
│ WorkerAddress │
│ (MessagePack map) │
│ │
│ tcp -> bytes │
│ grpc -> bytes │
│ ucx -> bytes │
└────────────────────────┘
```
When a Velo client receives a `PeerInfo`, it can:
1. Check `available_transports()` to see what's supported
2. Extract the relevant endpoint via `get_entry(key)`
3. Register the peer with its own transports
This design decouples the common types from specific transport implementations.
## Wire Format
- **InstanceId**: Serializes as a UUID string (JSON) or 16 bytes (binary)
- **WorkerId**: Serializes as a u64
- **WorkerAddress**: Serializes as a byte array (the MessagePack-encoded map)
- **PeerInfo**: Serializes as a struct with `instance_id` and `worker_address` fields
All types implement `serde::Serialize` and `serde::Deserialize`.
lib/velo-common/src/address.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Address types for peer discovery.
//!
//! This module provides types for representing worker addresses and peer information:
//! - [`WorkerAddress`]: Opaque byte representation of a peer's network address
//! - [`PeerInfo`]: Combined instance ID and worker address for a discovered peer
//!
//! These types are intentionally transport-agnostic, storing addresses as opaque bytes.
//! The interpretation of these bytes is left to the active message runtime.
use crate::identity::{InstanceId, WorkerId};
use crate::transport::TransportKey;
use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use xxhash_rust::xxh3::xxh3_64;
/// Errors that can occur when working with WorkerAddress.
#[derive(Debug, thiserror::Error)]
pub enum WorkerAddressError {
/// Attempted to add a key that already exists
#[error("Key already exists: {0}")]
KeyExists(String),
/// Attempted to access or remove a key that doesn't exist
#[error("Key not found: {0}")]
KeyNotFound(String),
/// Failed to encode the map to bytes
#[error("Encoding error: {0}")]
EncodingError(#[from] rmp_serde::encode::Error),
/// Failed to decode bytes to map
#[error("Decoding error: {0}")]
DecodingError(#[from] rmp_serde::decode::Error),
/// Encountered an unsupported format version
#[error("Unsupported format version: {0}")]
UnsupportedVersion(u8),
/// The data format is invalid
#[error("Invalid format: {0}")]
InvalidFormat(String),
}
/// Opaque worker address for discovery.
///
/// This is a transport-agnostic representation of a peer's network address.
/// The bytes are opaque to discovery and are interpreted by the active message runtime.
///
/// # Checksum
///
/// WorkerAddress implements a checksum via xxh3_64 for quick comparison during
/// re-registration validation.
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct WorkerAddress(Bytes);
// Custom Serialize/Deserialize to handle Bytes
impl Serialize for WorkerAddress {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serde_bytes::serialize(self.0.as_ref(), serializer)
}
}
impl<'de> Deserialize<'de> for WorkerAddress {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let bytes: Vec<u8> = serde_bytes::deserialize(deserializer)?;
Ok(WorkerAddress(Bytes::from(bytes)))
}
}
impl WorkerAddress {
/// Create a WorkerAddress from pre-encoded bytes.
///
/// This is used by transport implementations to construct addresses from
/// MessagePack-encoded map data. The bytes are assumed to be valid MessagePack.
pub fn from_encoded(bytes: impl Into<Bytes>) -> Self {
Self(bytes.into())
}
/// Get the underlying bytes.
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
/// Get the bytes as a Bytes object.
pub fn to_bytes(&self) -> Bytes {
self.0.clone()
}
/// Compute a checksum of this address for validation.
///
/// This is used to quickly check if an address has changed during re-registration.
pub fn checksum(&self) -> u64 {
xxh3_64(self.as_bytes())
}
/// Get the list of available transport keys in this address.
///
/// Returns the keys from the internal map as `TransportKey` for type-safe efficient
/// storage and sharing. This allows callers to see what transport types or endpoints
/// are available without exposing the full map.
///
/// # Errors
///
/// Returns an error if the internal bytes cannot be decoded as a valid MessagePack map.
///
/// # Example
///
/// ```no_run
/// # use velo_common::{WorkerAddress, TransportKey};
/// # let address: WorkerAddress = unimplemented!();
/// let transports = address.available_transports().unwrap();
/// if transports.contains(&TransportKey::from("tcp")) {
/// // TCP transport is available
/// }
/// ```
pub fn available_transports(&self) -> Result<Vec<TransportKey>, WorkerAddressError> {
let map = decode_to_map(self.as_bytes())?;
Ok(map.keys().cloned().map(TransportKey::from).collect())
}
/// Get a single entry from the internal map.
///
/// This decodes the address and extracts the entry for the given key.
///
/// Accepts any type that can be converted to a string reference, including
/// `&str`, `String`, `&String`, and `TransportKey`.
///
/// # Errors
///
/// Returns an error if the internal bytes cannot be decoded as a valid MessagePack map.
pub fn get_entry(&self, key: impl AsRef<str>) -> Result<Option<Bytes>, WorkerAddressError> {
let map = decode_to_map(self.as_bytes())?;
Ok(map.get(key.as_ref()).cloned())
}
}
impl fmt::Debug for WorkerAddress {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("WorkerAddress")
.field(&format_args!(
"len={}, xxh3_64=0x{:016x}",
self.0.len(),
self.checksum()
))
.finish()
}
}
impl fmt::Display for WorkerAddress {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "WorkerAddress(xxh3_64=0x{:016x})", self.checksum())
}
}
// ============================================================================
// Internal Decoding Helper
// ============================================================================
/// Decode WorkerAddress bytes from MessagePack into a map.
fn decode_to_map(bytes: &[u8]) -> Result<HashMap<Arc<str>, Bytes>, WorkerAddressError> {
if bytes.is_empty() {
return Err(WorkerAddressError::InvalidFormat("Empty bytes".to_string()));
}
// Decode MessagePack
let decoded: HashMap<String, Vec<u8>> = rmp_serde::from_slice(bytes)?;
// Convert to HashMap<Arc<str>, Bytes>
Ok(decoded
.into_iter()
.map(|(k, v)| (Arc::from(k.as_str()), Bytes::from(v)))
.collect())
}
/// Peer information combining instance ID and worker address.
///
/// This is the primary type returned by discovery lookups. It contains everything
/// needed to connect to and identify a peer.
///
/// # Example
///
/// ```no_run
/// # // WorkerAddress is created internally, this is simplified for docs
/// use velo_common::{InstanceId, PeerInfo};
/// # use velo_common::WorkerAddress;
/// # let address: WorkerAddress = unimplemented!();
///
/// let instance_id = InstanceId::new_v4();
/// let peer_info = PeerInfo::new(instance_id, address);
///
/// assert_eq!(peer_info.instance_id(), instance_id);
/// assert_eq!(peer_info.worker_id(), instance_id.worker_id());
/// ```
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct PeerInfo {
/// The instance ID of the peer
pub instance_id: InstanceId,
/// The worker address for connecting to the peer
pub worker_address: WorkerAddress,
}
impl PeerInfo {
/// Create a new PeerInfo.
pub fn new(instance_id: InstanceId, worker_address: WorkerAddress) -> Self {
Self {
instance_id,
worker_address,
}
}
/// Get the instance ID.
pub fn instance_id(&self) -> InstanceId {
self.instance_id
}
/// Get the worker ID (derived from instance ID).
pub fn worker_id(&self) -> WorkerId {
self.instance_id.worker_id()
}
/// Get a reference to the worker address.
pub fn worker_address(&self) -> &WorkerAddress {
&self.worker_address
}
/// Get the worker address checksum for validation.
pub fn address_checksum(&self) -> u64 {
self.worker_address.checksum()
}
/// Consume self and return the worker address.
pub fn into_address(self) -> WorkerAddress {
self.worker_address
}
/// Decompose into instance ID and worker address.
pub fn into_parts(self) -> (InstanceId, WorkerAddress) {
(self.instance_id, self.worker_address)
}
}
#[cfg(test)]
mod tests {
use super::*;
// Helper to create a test address with MessagePack encoding
fn make_test_address(entries: &[(&str, &[u8])]) -> WorkerAddress {
let map: HashMap<String, Vec<u8>> = entries
.iter()
.map(|(k, v)| (k.to_string(), v.to_vec()))
.collect();
let encoded = rmp_serde::to_vec(&map).unwrap();
WorkerAddress::from_encoded(encoded)
}
#[test]
fn test_worker_address_from_encoded() {
let address = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
// Verify we can get the entry back
let entry = address.get_entry("endpoint").unwrap();
assert_eq!(entry, Some(Bytes::from_static(b"tcp://127.0.0.1:5555")));
}
#[test]
fn test_worker_address_checksum() {
let address1 = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let address2 = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let address3 = make_test_address(&[("endpoint", b"tcp://127.0.0.1:6666")]);
// Same content = same checksum
assert_eq!(address1.checksum(), address2.checksum());
// Different content = different checksum
assert_ne!(address1.checksum(), address3.checksum());
}
#[test]
fn test_worker_address_equality() {
let address1 = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let address2 = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let address3 = make_test_address(&[("endpoint", b"tcp://127.0.0.1:6666")]);
assert_eq!(address1, address2);
assert_ne!(address1, address3);
}
#[test]
fn test_worker_address_debug() {
let address = make_test_address(&[("test", b"value")]);
let debug_str = format!("{:?}", address);
assert!(debug_str.contains("WorkerAddress"));
assert!(debug_str.contains("len="));
assert!(debug_str.contains("xxh3_64="));
}
#[test]
fn test_available_transports() {
let address = make_test_address(&[
("tcp", b"tcp://127.0.0.1:5555"),
("rdma", b"rdma://10.0.0.1:6666"),
("udp", b"udp://127.0.0.1:7777"),
]);
let transports = address.available_transports().unwrap();
assert_eq!(transports.len(), 3);
assert!(transports.contains(&TransportKey::from("tcp")));
assert!(transports.contains(&TransportKey::from("rdma")));
assert!(transports.contains(&TransportKey::from("udp")));
}
#[test]
fn test_available_transports_empty() {
let address = make_test_address(&[]);
let transports = address.available_transports().unwrap();
assert_eq!(transports.len(), 0);
}
#[test]
fn test_get_entry() {
let address =
make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555"), ("protocol", b"tcp")]);
// Get existing entry
assert_eq!(
address.get_entry("endpoint").unwrap().unwrap(),
Bytes::from_static(b"tcp://127.0.0.1:5555")
);
// Get nonexistent entry
assert!(address.get_entry("nonexistent").unwrap().is_none());
}
#[test]
fn test_get_entry_with_transport_key() {
let address = make_test_address(&[
("tcp", b"tcp://127.0.0.1:5555"),
("rdma", b"rdma://10.0.0.1:6666"),
]);
// Test get_entry with TransportKey
let tcp_key = TransportKey::from("tcp");
let result = address.get_entry(tcp_key).unwrap();
assert_eq!(result, Some(Bytes::from_static(b"tcp://127.0.0.1:5555")));
// Test get_entry with String
let result = address.get_entry(String::from("rdma")).unwrap();
assert_eq!(result, Some(Bytes::from_static(b"rdma://10.0.0.1:6666")));
}
#[test]
fn test_peer_info_creation() {
let instance_id = InstanceId::new_v4();
let address = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let peer_info = PeerInfo::new(instance_id, address.clone());
assert_eq!(peer_info.instance_id(), instance_id);
assert_eq!(peer_info.worker_id(), instance_id.worker_id());
assert_eq!(peer_info.worker_address(), &address);
}
#[test]
fn test_peer_info_checksum() {
let instance_id = InstanceId::new_v4();
let address = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let peer_info = PeerInfo::new(instance_id, address.clone());
assert_eq!(peer_info.address_checksum(), address.checksum());
}
#[test]
fn test_peer_info_into_address() {
let instance_id = InstanceId::new_v4();
let address = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let peer_info = PeerInfo::new(instance_id, address.clone());
let extracted_address = peer_info.into_address();
assert_eq!(extracted_address, address);
}
#[test]
fn test_peer_info_into_parts() {
let instance_id = InstanceId::new_v4();
let address = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let peer_info = PeerInfo::new(instance_id, address.clone());
let (extracted_id, extracted_address) = peer_info.into_parts();
assert_eq!(extracted_id, instance_id);
assert_eq!(extracted_address, address);
}
#[test]
fn test_peer_info_serde() {
let instance_id = InstanceId::new_v4();
let address = make_test_address(&[("endpoint", b"tcp://127.0.0.1:5555")]);
let peer_info = PeerInfo::new(instance_id, address);
// Serialize to JSON
let json = serde_json::to_string(&peer_info).unwrap();
// Deserialize back
let deserialized: PeerInfo = serde_json::from_str(&json).unwrap();
assert_eq!(deserialized.instance_id(), instance_id);
assert_eq!(deserialized.worker_id(), instance_id.worker_id());
// Verify the entry is preserved
let entry = deserialized.worker_address().get_entry("endpoint").unwrap();
assert_eq!(entry, Some(Bytes::from_static(b"tcp://127.0.0.1:5555")));
}
}
lib/velo-common/src/identity.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Identity types for the active message system.
//!
//! This module provides strongly-typed wrappers for instance and worker identifiers:
//! - [`InstanceId`]: Unique runtime instance identifier (wraps UUID)
//! - [`WorkerId`]: Deterministic 64-bit worker identifier derived from InstanceId
//!
//! # Design Principles
//!
//! 1. **Type Safety**: InstanceId cannot be confused with message IDs or other UUIDs
//! 2. **Deterministic Derivation**: WorkerId is always computed from InstanceId (xxh3_64 hash)
//! 3. **Single Source of Truth**: InstanceId is the primary identifier, WorkerId is derived
use serde::{Deserialize, Serialize};
use std::fmt;
use uuid::Uuid;
use xxhash_rust::xxh3::xxh3_64;
/// Unique identifier for a runtime instance.
///
/// This is a UUID-based identifier that uniquely identifies a running instance
/// of the active message runtime. It is used for:
/// - Transport-level addressing
/// - Discovery registration
/// - Routing table management
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct InstanceId(Uuid);
impl InstanceId {
/// Create a new random v4 InstanceId.
///
/// This is exposed for testing and special cases. In production, use
/// [`InstanceFactory::create()`] instead.
pub fn new_v4() -> Self {
loop {
let instance_id = InstanceId(Uuid::new_v4());
let worker_id = WorkerId::from(&instance_id);
if worker_id.as_u64() != 0 {
return instance_id;
}
}
}
/// Derive the deterministic WorkerId from this InstanceId.
///
/// WorkerId is computed using xxh3_64 hash of the UUID bytes.
/// This ensures a 1:1 mapping between InstanceId and WorkerId.
pub fn worker_id(&self) -> WorkerId {
WorkerId::from(self)
}
/// Get a reference to the underlying UUID.
pub fn as_uuid(&self) -> &Uuid {
&self.0
}
/// Get the underlying UUID as a u128.
pub fn as_u128(&self) -> u128 {
self.0.as_u128()
}
/// Get the underlying UUID as bytes.
pub fn as_bytes(&self) -> &[u8; 16] {
self.0.as_bytes()
}
}
impl fmt::Display for InstanceId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<Uuid> for InstanceId {
fn from(uuid: Uuid) -> Self {
Self(uuid)
}
}
impl From<InstanceId> for Uuid {
fn from(id: InstanceId) -> Self {
id.0
}
}
impl AsRef<Uuid> for InstanceId {
fn as_ref(&self) -> &Uuid {
&self.0
}
}
/// Deterministic 64-bit worker identifier derived from InstanceId.
///
/// WorkerId enables embedding instance identity into fixed-size handles that can be
/// passed with value semantics. A `u128` is the largest integer that can be passed
/// by value, making it ideal for handles that encode both routing and event information.
///
/// WorkerId is used in:
/// - `EventHandle` (velo): Uses 64 bits for WorkerId + 64 bits for event details
/// - `EventRoutingTable` (velo): Maps worker_id → instance_id for event routing
/// - Discovery systems: Lookup key for peer information
///
/// WorkerId is **always derived** from InstanceId using xxh3_64 hash.
/// This ensures consistency across the system without needing to store both values.
///
/// # Example
///
/// ```ignore
/// use velo_common::{InstanceId, WorkerId};
///
/// # fn get_instance_id() -> InstanceId { unimplemented!() }
/// let instance_id = get_instance_id(); // From ActiveMessageClient
/// let worker_id = instance_id.worker_id();
///
/// // WorkerId is deterministic
/// assert_eq!(worker_id, instance_id.worker_id());
/// ```
#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, PartialOrd, Ord)]
#[serde(transparent)]
pub struct WorkerId(u64);
impl WorkerId {
/// Create a WorkerId from a raw u64 value.
///
/// This is used when decoding WorkerIds from event handles or wire formats.
/// External users should always derive WorkerId via `instance_id.worker_id()`.
pub fn from_u64(value: u64) -> Self {
Self(value)
}
/// Get the underlying u64 value.
#[inline(always)]
pub fn as_u64(&self) -> u64 {
self.0
}
}
impl fmt::Display for WorkerId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<&InstanceId> for WorkerId {
/// Derive WorkerId from InstanceId using xxh3_64 hash.
///
/// This is the canonical way to compute WorkerId - it should never be
/// constructed any other way to ensure consistency.
fn from(id: &InstanceId) -> Self {
Self(xxh3_64(id.as_uuid().as_bytes()))
}
}
impl From<InstanceId> for WorkerId {
fn from(id: InstanceId) -> Self {
Self::from(&id)
}
}
impl From<WorkerId> for u64 {
fn from(id: WorkerId) -> Self {
id.0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_instance_id_creation() {
let id1 = InstanceId::new_v4();
let id2 = InstanceId::new_v4();
// Different instances have different IDs
assert_ne!(id1, id2);
// Can convert to/from UUID
let uuid: Uuid = id1.into();
let id3 = InstanceId::from(uuid);
assert_eq!(id1, id3);
}
#[test]
fn test_worker_id_deterministic() {
let instance_id = InstanceId::new_v4();
// WorkerId is deterministic
let worker_id1 = instance_id.worker_id();
let worker_id2 = instance_id.worker_id();
assert_eq!(worker_id1, worker_id2);
// Different instances have different worker IDs
let other_instance = InstanceId::new_v4();
let other_worker = other_instance.worker_id();
assert_ne!(worker_id1, other_worker);
}
#[test]
fn test_worker_id_from_conversion() {
let instance_id = InstanceId::new_v4();
// Both From implementations work
let worker_id1 = WorkerId::from(&instance_id);
let worker_id2 = WorkerId::from(instance_id);
assert_eq!(worker_id1, worker_id2);
// Matches .worker_id() method
assert_eq!(worker_id1, instance_id.worker_id());
}
#[test]
fn test_instance_id_display() {
let instance_id = InstanceId::new_v4();
let display = format!("{}", instance_id);
let uuid_display = format!("{}", instance_id.as_uuid());
assert_eq!(display, uuid_display);
}
#[test]
fn test_worker_id_display() {
let instance_id = InstanceId::new_v4();
let worker_id = instance_id.worker_id();
let display = format!("{}", worker_id);
let u64_display = format!("{}", worker_id.as_u64());
assert_eq!(display, u64_display);
}
#[test]
fn test_instance_id_serde() {
let instance_id = InstanceId::new_v4();
// Serialize as JSON
let json = serde_json::to_string(&instance_id).unwrap();
// Should be a plain UUID string
let uuid_json = serde_json::to_string(instance_id.as_uuid()).unwrap();
assert_eq!(json, uuid_json);
// Deserialize back
let deserialized: InstanceId = serde_json::from_str(&json).unwrap();
assert_eq!(instance_id, deserialized);
}
#[test]
fn test_worker_id_serde() {
let worker_id = InstanceId::new_v4().worker_id();
// Serialize as JSON
let json = serde_json::to_string(&worker_id).unwrap();
// Should be a plain u64
let u64_json = serde_json::to_string(&worker_id.as_u64()).unwrap();
assert_eq!(json, u64_json);
// Deserialize back
let deserialized: WorkerId = serde_json::from_str(&json).unwrap();
assert_eq!(worker_id, deserialized);
}
#[test]
fn test_worker_id_u64_conversion() {
let instance_id = InstanceId::new_v4();
let worker_id = instance_id.worker_id();
let raw_u64 = worker_id.as_u64();
let reconstructed = WorkerId::from_u64(raw_u64);
assert_eq!(worker_id, reconstructed);
}
}
lib/velo-common/src/lib.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
#![doc = include_str!("../README.md")]
mod address;
mod identity;
mod transport;
// Re-export all public types
pub use address::{PeerInfo, WorkerAddress, WorkerAddressError};
pub use identity::{InstanceId, WorkerId};
pub use transport::TransportKey;
lib/velo-common/src/transport.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Transport key type for type-safe transport identification.
use std::fmt;
use std::sync::Arc;
/// A type-safe wrapper around transport keys for WorkerAddress.
///
/// This provides a zero-cost abstraction over `Arc<str>` with type safety
/// to prevent accidentally mixing transport keys with other string types.
///
/// # Examples
///
/// ```
/// use velo_common::TransportKey;
///
/// let key = TransportKey::new("tcp");
/// assert_eq!(key.as_str(), "tcp");
///
/// // Ergonomic conversions
/// let key2: TransportKey = "rdma".into();
/// let key3 = TransportKey::from("udp");
///
/// // Use in collections
/// use std::collections::HashMap;
/// let mut transports = HashMap::new();
/// transports.insert(TransportKey::from("tcp"), "tcp://127.0.0.1:5555");
/// ```
#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct TransportKey(Arc<str>);
impl TransportKey {
/// Create a new TransportKey from any type that can be converted into Arc<str>.
pub fn new(key: impl Into<Arc<str>>) -> Self {
Self(key.into())
}
/// Get the key as a string slice.
pub fn as_str(&self) -> &str {
&self.0
}
}
// Deref to str for ergonomic usage
impl std::ops::Deref for TransportKey {
type Target = str;
fn deref(&self) -> &Self::Target {
&self.0
}
}
// AsRef for flexible parameter types
impl AsRef<str> for TransportKey {
fn as_ref(&self) -> &str {
&self.0
}
}
// From conversions for ergonomic construction
impl From<&str> for TransportKey {
fn from(s: &str) -> Self {
Self(Arc::from(s))
}
}
impl From<String> for TransportKey {
fn from(s: String) -> Self {
Self(Arc::from(s))
}
}
impl From<Arc<str>> for TransportKey {
fn from(s: Arc<str>) -> Self {
Self(s)
}
}
impl From<&String> for TransportKey {
fn from(s: &String) -> Self {
Self(Arc::from(s.as_str()))
}
}
impl From<TransportKey> for String {
fn from(val: TransportKey) -> Self {
val.0.to_string()
}
}
// Borrow trait for HashMap lookups with &str
impl std::borrow::Borrow<str> for TransportKey {
fn borrow(&self) -> &str {
&self.0
}
}
// Display for printing
impl fmt::Display for TransportKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::{HashMap, HashSet};
#[test]
fn test_transport_key_creation() {
// Test new() method
let key1 = TransportKey::new("tcp");
assert_eq!(key1.as_str(), "tcp");
// Test From<&str>
let key2: TransportKey = "rdma".into();
assert_eq!(key2.as_str(), "rdma");
// Test From<String>
let key3 = TransportKey::from(String::from("udp"));
assert_eq!(key3.as_str(), "udp");
// Test From<&String>
let s = String::from("grpc");
let key4 = TransportKey::from(&s);
assert_eq!(key4.as_str(), "grpc");
// Test From<Arc<str>>
let arc_str: Arc<str> = Arc::from("http");
let key5 = TransportKey::from(arc_str);
assert_eq!(key5.as_str(), "http");
}
#[test]
fn test_transport_key_deref() {
let key = TransportKey::from("tcp");
// Deref to str methods should work
assert_eq!(key.len(), 3);
assert_eq!(key.chars().count(), 3);
assert!(key.starts_with("tc"));
assert!(key.ends_with("cp"));
// Can use str slicing through Deref
assert_eq!(&key[0..2], "tc");
}
#[test]
fn test_transport_key_as_ref() {
let key = TransportKey::from("tcp");
// AsRef<str> allows passing to functions expecting &str
fn takes_str_ref(s: &str) -> usize {
s.len()
}
assert_eq!(takes_str_ref(&key), 3);
assert_eq!(takes_str_ref(key.as_ref()), 3);
}
#[test]
fn test_transport_key_display() {
let key = TransportKey::from("tcp");
assert_eq!(format!("{}", key), "tcp");
assert_eq!(key.to_string(), "tcp");
}
#[test]
fn test_transport_key_debug() {
let key = TransportKey::from("tcp");
let debug_str = format!("{:?}", key);
assert!(debug_str.contains("TransportKey"));
assert!(debug_str.contains("tcp"));
}
#[test]
fn test_transport_key_equality() {
let key1 = TransportKey::from("tcp");
let key2 = TransportKey::from("tcp");
let key3 = TransportKey::from("rdma");
assert_eq!(key1, key2);
assert_ne!(key1, key3);
// Test with different source types
let key4: TransportKey = String::from("tcp").into();
assert_eq!(key1, key4);
}
#[test]
fn test_transport_key_ordering() {
let mut keys = [
TransportKey::from("udp"),
TransportKey::from("tcp"),
TransportKey::from("rdma"),
TransportKey::from("grpc"),
];
keys.sort();
assert_eq!(keys[0], TransportKey::from("grpc"));
assert_eq!(keys[1], TransportKey::from("rdma"));
assert_eq!(keys[2], TransportKey::from("tcp"));
assert_eq!(keys[3], TransportKey::from("udp"));
}
#[test]
fn test_transport_key_hash() {
let mut set = HashSet::new();
set.insert(TransportKey::from("tcp"));
set.insert(TransportKey::from("rdma"));
set.insert(TransportKey::from("tcp")); // Duplicate
assert_eq!(set.len(), 2);
assert!(set.contains(&TransportKey::from("tcp")));
assert!(set.contains(&TransportKey::from("rdma")));
assert!(!set.contains(&TransportKey::from("udp")));
}
#[test]
fn test_transport_key_in_hashmap() {
let mut map = HashMap::new();
map.insert(TransportKey::from("tcp"), "tcp://127.0.0.1:5555");
map.insert(TransportKey::from("rdma"), "rdma://10.0.0.1:6666");
// Can lookup with TransportKey
assert_eq!(
map.get(&TransportKey::from("tcp")),
Some(&"tcp://127.0.0.1:5555")
);
// Can lookup with &str via Borrow trait
assert_eq!(map.get("tcp"), Some(&"tcp://127.0.0.1:5555"));
assert_eq!(map.get("rdma"), Some(&"rdma://10.0.0.1:6666"));
assert_eq!(map.get("udp"), None);
}
#[test]
fn test_transport_key_clone() {
let key1 = TransportKey::from("tcp");
let key2 = key1.clone();
assert_eq!(key1, key2);
assert_eq!(key1.as_str(), key2.as_str());
// Verify Arc is shared (same pointer)
let ptr1 = key1.as_str().as_ptr();
let ptr2 = key2.as_str().as_ptr();
assert_eq!(ptr1, ptr2);
}
}
lib/velo-transports/Cargo.toml 0 → 100644
View file @ 63d7c01c
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
[package]
name = "velo-transports"
version.workspace = true
edition.workspace = true
description.workspace = true
authors.workspace = true
license.workspace = true
homepage.workspace = true
repository.workspace = true
keywords.workspace = true
[features]
default = ["http", "nats", "grpc"]
ucx = ["dep:async-ucx"]
http = ["dep:axum", "dep:reqwest"]
nats = ["dep:async-nats", "dep:bs58"]
grpc = ["dep:tonic", "dep:prost", "dep:tower", "dep:hyper", "dep:http", "dep:http-body", "dep:http-body-util", "dep:hyper-util", "dep:tokio-stream"]
[dependencies]
velo-common = { workspace = true }
anyhow = { workspace = true }
bytes = { workspace = true }
dashmap = { workspace = true }
# derive_builder = { workspace = true }
parking_lot.workspace = true
# serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true }
tracing = { workspace = true }
# uuid = { workspace = true, features = ["serde", "v4"] }
# xxhash-rust = { workspace = true }
# base64 = "0.22"
flume = "0.12.0"
futures = "0.3"
# hashbrown = "0.16"
# lru = { version = "0.16", features = ["hashbrown"]}
rmp-serde = "1.1"
# serde_bytes = "0.11"
socket2 = "0.6"
tokio-util = { version = "0.7", features = ["codec"] }
[target.'cfg(target_os = "linux")'.dependencies]
nix = { version = "0.30", features = ["sched"] }
# optional dependencies
async-ucx = { git = "https://github.com/ryanolson/async-ucx.git", branch = "ryan/ucx", optional = true, features = ["am", "event"] }
axum = { version = "0.8", optional = true }
reqwest = { workspace = true , optional = true }
async-nats = { workspace = true, optional = true }
bs58 = { version = "0.5", optional = true }
prost = { version = "0.13", optional = true }
tonic = { version = "0.13.1", optional = true }
tower = { version = "0.5", optional = true }
hyper = { version = "1.0", optional = true }
http = { version = "1.0", optional = true }
http-body = { version = "1.0", optional = true }
http-body-util = { version = "0.1", optional = true }
hyper-util = { version = "0.1", optional = true, features = ["tokio", "server", "server-auto"] }
tokio-stream = { version = "0.1", optional = true, features = ["sync"] }
[dev-dependencies]
tokio = { workspace = true, features = ["test-util", "macros"] }
tower = "0.5"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
[build-dependencies]
tonic-build = "0.13.1"
lib/velo-transports/README.md 0 → 100644
View file @ 63d7c01c
# velo-transports
Multi-transport active message routing for distributed systems.
## Overview
`velo-transports` abstracts TCP, HTTP, NATS, gRPC, and UCX behind a unified `Transport` trait. It provides:
- **Zero-copy `Bytes`** — inbound frames share the receive buffer via `bytes::Bytes` slicing
- **Fire-and-forget with error callbacks** — `send_message` is non-blocking; failures are reported through `TransportErrorHandler`
- **Priority-based peer routing** — each peer is registered with all compatible transports; the highest-priority one becomes the primary
- **3-phase graceful shutdown** — Gate → Drain → Teardown with RAII in-flight tracking
- **Peer discovery via `WorkerAddress`** — MessagePack-encoded map of transport key → endpoint bytes
## Architecture
```mermaid
graph TD
VB[VeloBackend] -->|priority routing| TCP[TcpTransport]
VB -->|priority routing| HTTP[HttpTransport]
VB -->|priority routing| NATS[NatsTransport]
VB -->|priority routing| GRPC[GrpcTransport]
VB -->|priority routing| UCX[UcxTransport]
TCP -->|inbound| DS[DataStreams]
HTTP -->|inbound| DS
NATS -->|inbound| DS
GRPC -->|inbound| DS
UCX -->|inbound| DS
DS -->|message_stream| APP[Application]
DS -->|response_stream| APP
DS -->|event_stream| APP
```
`VeloBackend` is the central orchestrator. It starts each transport, builds a composite `WorkerAddress` advertising all endpoints, and manages the peer registry. Inbound frames arrive via three independent `flume` channels in `DataStreams`.
## Feature Flags
| Feature | Default | Dependencies | Description |
|---------|---------|-------------|-------------|
| `http` | ✓ | axum, reqwest | HTTP transport with Axum server |
| `nats` | ✓ | async-nats, bs58 | NATS pub-sub transport |
| `grpc` | ✓ | tonic, prost, tower | gRPC bidirectional streaming |
| `ucx` | ✗ | async-ucx | UCX/RDMA transport (requires UCX libraries) |
TCP transport is always available (no feature gate).
## Transport Summary
| Transport | Protocol | Framing | Key Properties |
|-----------|----------|---------|----------------|
| **TCP** | Raw TCP | 11-byte preamble + header + payload | Zero-copy codec, DashMap connection pool, CPU pinning, keepalive |
| **HTTP** | HTTP/1.1 POST | Base64 header in `X-Transport-Header`, raw body | Fire-and-forget (202 Accepted), Axum server, drain → 503 |
| **NATS** | NATS pub-sub | Base64 header in NATS HeaderMap, raw payload | Subject scheme `velo.{b58}.{type}`, request/reply health, drain via unsub |
| **gRPC** | HTTP/2 streaming | Protobuf `FramedData` wrapper (preamble + header + payload) | Bidirectional streaming, tonic channels, exponential backoff reconnect |
| **UCX** | RDMA Active Messages | 3 fixed lanes (msg/resp/event) | `LocalSet` thread for `Rc<Worker>`, lazy endpoints, zero-copy RDMA |
## Shutdown Model
Graceful shutdown follows three phases:
1. **Gate** — `begin_drain()` flips an atomic flag. Transports reject new inbound `Message` frames (TCP sends `ShuttingDown` response, HTTP returns 503, NATS unsubscribes from message subject).
2. **Drain** — `wait_for_drain()` blocks until all `InFlightGuard`s are dropped. Policy is either `WaitForever` or `Timeout(Duration)`.
3. **Teardown** — Cancel the teardown token, stopping all listener loops and writer tasks. Call `shutdown()` on each transport to clean up connections.
Response, Ack, and Event frames continue flowing during drain so in-flight work can complete.
## Wire Format (TCP)
```text
┌──────────────┬───────────┬──────────────┬───────────────┬────────┬─────────┐
│ version (2B) │ type (1B) │ hdr_len (4B) │ pay_len (4B) │ header │ payload │
│ u16 BE │ u8 │ u32 BE │ u32 BE │ bytes │ bytes │
└──────────────┴───────────┴──────────────┴───────────────┴────────┴─────────┘
```
- **version**: Schema version (currently 1)
- **type**: `Message(0)`, `Response(1)`, `Ack(2)`, `Event(3)`, `ShuttingDown(4)`
- **hdr_len / pay_len**: Lengths of the following header and payload sections
- **Max frame size**: 16 MB
The gRPC transport wraps the same preamble + header + payload in a Protobuf `FramedData` message.
lib/velo-transports/build.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
fn main() -> Result<(), Box<dyn std::error::Error>> {
// Only compile proto files if grpc feature is enabled
#[cfg(feature = "grpc")]
{
tonic_build::compile_protos("proto/velo.proto")?;
}
Ok(())
}
lib/velo-transports/docs/tcp.md 0 → 100644
View file @ 63d7c01c
# TCP Transport Design
## Overview
The TCP transport provides high-performance, zero-copy message delivery over raw TCP connections. It uses a custom 11-byte frame preamble for minimal overhead and supports CPU pinning for predictable latency.
## Connection Management
### DashMap Connection Pool
Peer connections are managed via two `DashMap` instances:
- `peers: DashMap<InstanceId, SocketAddr>` — registered peer addresses
- `connections: DashMap<InstanceId, ConnectionHandle>` — active writer task handles
Connections are established lazily on first `send_message()` call. Each connection spawns a dedicated writer task that owns the TCP stream.
### Writer Tasks
Each `ConnectionHandle` wraps a bounded `flume::Sender<SendTask>` (default capacity: 256). The send path:
1. **Fast path**: `try_send()` on existing connection — non-blocking, no allocation
2. **Slow path (full)**: `send_async()` via spawned task — applies backpressure
3. **Slow path (new)**: `get_or_create_connection()` — establishes TCP connection and spawns writer
The writer task loop:
```
recv_async(SendTask) → encode_frame(&mut stream, ...) → loop
```
## TcpFrameCodec
### Wire Format
```
[u16 BE: schema_version(1)] [u8: frame_type] [u32 BE: header_len] [u32 BE: payload_len] [header] [payload]
```
Total preamble: 11 bytes. Maximum frame: 16 MB.
### Decoder State Machine
The codec uses a two-state decoder for streaming TCP:
```
AwaitingHeader ──(11 bytes available)──→ AwaitingData ──(data available)──→ emit frame, reset
```
Zero-copy is achieved via `BytesMut::split_to().freeze()` — the output `Bytes` share the underlying receive buffer.
### Encoder
`encode_frame()` writes three segments via `write_all()`:
1. Preamble (11 bytes)
2. Header bytes
3. Payload bytes
`write_vectored()` is intentionally not used because it doesn't guarantee writing all bytes for payloads exceeding the kernel send buffer (~128KB).
## TCP Listener
### Frame Routing
Incoming frames are routed based on `MessageType`:
| MessageType | Target Stream |
|------------|---------------|
| Message | `message_stream` |
| Response | `response_stream` |
| Ack, Event | `event_stream` |
| ShuttingDown | `response_stream` (for correlation) |
### Drain Behavior
During drain (`ShutdownState::is_draining()`):
- **Message** frames are rejected: a `ShuttingDown` frame is sent back with the original header for correlation
- **Response/Ack/Event** frames pass through normally
### CPU Pinning (Linux)
`RuntimeConfig::CpuPin(cpu_id)` creates a single-threaded tokio runtime with the thread pinned to the specified CPU core via `nix::sched::sched_setaffinity`. This provides predictable latency by avoiding context switches.
On non-Linux platforms, `CpuPin` falls back to a regular single-threaded runtime with a warning.
## Socket Configuration
Both listener and writer sides configure:
- `TCP_NODELAY` — disable Nagle's algorithm for low-latency framing
- `SO_LINGER(1s)` — ensure clean socket shutdown
- `TCP_KEEPALIVE` — 60s idle time, 10s probe interval
- **Buffer sizes** — 1 MB send/receive buffers for high throughput
lib/velo-transports/proto/velo.proto 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
syntax = "proto3";
package velo.streaming.v1;
// Velo streaming service for active messages
//
// This service uses bidirectional streaming to send pre-framed messages.
// Messages are wrapped in FramedData which contains our custom TCP frame format.
service VeloStreaming
{
// Bidirectional streaming RPC
// Client sends framed messages, server receives them
// The response stream is unused (empty) in our implementation
rpc Stream(stream FramedData) returns (stream FramedData);
}
// Wrapper for pre-framed message data
//
// The fields map to TCP frame segments:
// preamble: [u16: version][u8: msg_type][u32: header_len][u32: payload_len]
// header: [header bytes]
// payload: [payload bytes]
message FramedData
{
bytes preamble = 1;
bytes header = 2;
bytes payload = 3;
}
lib/velo-transports/src/address.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Internal address builder for constructing WorkerAddress instances.
//!
//! This module provides the builder pattern for creating WorkerAddress instances
//! from transport-specific endpoint data. It is internal to velo-transports.
use bytes::Bytes;
use std::collections::HashMap;
use std::sync::Arc;
use velo_common::{WorkerAddress, WorkerAddressError};
/// Builder for constructing WorkerAddress instances.
///
/// This provides a mutable interface for collecting transport endpoints
/// before encoding them into the immutable WorkerAddress format.
#[derive(Debug, Clone, Default)]
pub(crate) struct WorkerAddressBuilder {
entries: HashMap<String, Bytes>,
}
impl WorkerAddressBuilder {
/// Create a new empty builder.
pub fn new() -> Self {
Self {
entries: HashMap::new(),
}
}
/// Add a new entry to the map.
///
/// Returns an error if the key already exists.
pub fn add_entry(
&mut self,
key: impl Into<String>,
value: impl Into<Bytes>,
) -> Result<(), WorkerAddressError> {
let key = key.into();
if self.entries.contains_key(&key) {
return Err(WorkerAddressError::KeyExists(key));
}
self.entries.insert(key, value.into());
Ok(())
}
/// Check if a key exists in the map.
#[allow(dead_code)]
pub fn has_entry(&self, key: &str) -> bool {
self.entries.contains_key(key)
}
/// Get a reference to an entry's value.
#[allow(dead_code)]
pub fn get_entry(&self, key: &str) -> Option<&Bytes> {
self.entries.get(key)
}
/// Merge another WorkerAddress into this builder.
///
/// This decodes the other address and attempts to add all its entries to this builder.
/// If any key from the other address already exists in this builder, returns an error
/// and leaves the builder unchanged.
pub fn merge(&mut self, other: &WorkerAddress) -> Result<(), WorkerAddressError> {
let map = decode_to_map(other.as_bytes())?;
// First check if any keys would conflict
for key in map.keys() {
if self.entries.contains_key(key.as_ref()) {
return Err(WorkerAddressError::KeyExists(key.to_string()));
}
}
// All keys are unique, now add them
for (key, value) in map {
self.entries.insert(key.to_string(), value);
}
Ok(())
}
/// Build the WorkerAddress from this builder.
///
/// This encodes the map into MessagePack binary format.
pub fn build(self) -> Result<WorkerAddress, WorkerAddressError> {
// Convert HashMap<String, Bytes> to HashMap<String, Vec<u8>> for MessagePack
let serializable: HashMap<String, Vec<u8>> = self
.entries
.into_iter()
.map(|(k, v)| (k, v.to_vec()))
.collect();
// Encode to MessagePack
let encoded = rmp_serde::to_vec(&serializable)?;
Ok(WorkerAddress::from_encoded(encoded))
}
}
/// Decode WorkerAddress bytes from MessagePack into a map.
fn decode_to_map(bytes: &[u8]) -> Result<HashMap<Arc<str>, Bytes>, WorkerAddressError> {
if bytes.is_empty() {
return Err(WorkerAddressError::InvalidFormat("Empty bytes".to_string()));
}
// Decode MessagePack
let decoded: HashMap<String, Vec<u8>> = rmp_serde::from_slice(bytes)?;
// Convert to HashMap<Arc<str>, Bytes>
Ok(decoded
.into_iter()
.map(|(k, v)| (Arc::from(k.as_str()), Bytes::from(v)))
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_builder_basic() {
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry("endpoint", Bytes::from_static(b"tcp://127.0.0.1:5555"))
.unwrap();
builder
.add_entry("protocol", Bytes::from_static(b"tcp"))
.unwrap();
assert!(builder.has_entry("endpoint"));
assert!(builder.has_entry("protocol"));
assert!(!builder.has_entry("nonexistent"));
let address = builder.build().unwrap();
assert!(!address.as_bytes().is_empty());
// Verify we can read the entries back
let entry = address.get_entry("endpoint").unwrap();
assert_eq!(entry, Some(Bytes::from_static(b"tcp://127.0.0.1:5555")));
}
#[test]
fn test_builder_add_duplicate_key() {
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry("key", Bytes::from_static(b"value1"))
.unwrap();
let result = builder.add_entry("key", Bytes::from_static(b"value2"));
assert!(matches!(result, Err(WorkerAddressError::KeyExists(_))));
}
#[test]
fn test_builder_merge() {
// Build first address
let mut builder1 = WorkerAddressBuilder::new();
builder1
.add_entry("tcp", Bytes::from_static(b"tcp://127.0.0.1:5555"))
.unwrap();
let address1 = builder1.build().unwrap();
// Build second address
let mut builder2 = WorkerAddressBuilder::new();
builder2
.add_entry("rdma", Bytes::from_static(b"rdma://10.0.0.1:6666"))
.unwrap();
let address2 = builder2.build().unwrap();
// Merge both into a new builder
let mut builder3 = WorkerAddressBuilder::new();
builder3.merge(&address1).unwrap();
builder3.merge(&address2).unwrap();
let final_address = builder3.build().unwrap();
// Verify both entries are present
assert_eq!(
final_address.get_entry("tcp").unwrap(),
Some(Bytes::from_static(b"tcp://127.0.0.1:5555"))
);
assert_eq!(
final_address.get_entry("rdma").unwrap(),
Some(Bytes::from_static(b"rdma://10.0.0.1:6666"))
);
}
#[test]
fn test_builder_merge_with_conflict() {
let mut builder1 = WorkerAddressBuilder::new();
builder1
.add_entry("tcp", Bytes::from_static(b"tcp://127.0.0.1:5555"))
.unwrap();
let address1 = builder1.build().unwrap();
let mut builder2 = WorkerAddressBuilder::new();
builder2
.add_entry("tcp", Bytes::from_static(b"tcp://different:5555"))
.unwrap();
let address2 = builder2.build().unwrap();
// Merge first address
let mut builder3 = WorkerAddressBuilder::new();
builder3.merge(&address1).unwrap();
// Try to merge conflicting address - should fail
let result = builder3.merge(&address2);
assert!(matches!(result, Err(WorkerAddressError::KeyExists(_))));
// Builder should be unchanged
assert!(builder3.has_entry("tcp"));
assert_eq!(
builder3.get_entry("tcp").unwrap(),
&Bytes::from_static(b"tcp://127.0.0.1:5555")
);
}
#[test]
fn test_empty_builder() {
let builder = WorkerAddressBuilder::new();
let address = builder.build().unwrap();
// Empty address should still be valid
let transports = address.available_transports().unwrap();
assert_eq!(transports.len(), 0);
}
// ========================================================================
// Integration tests: Verify WorkerAddressBuilder (velo-transports) produces
// addresses that WorkerAddress (velo-common) can correctly decode.
// These tests ensure the two crates stay in sync on the wire format.
// ========================================================================
#[test]
fn test_builder_address_integration_get_entry() {
// Build an address with multiple entries
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry("tcp", Bytes::from_static(b"tcp://127.0.0.1:5555"))
.unwrap();
builder
.add_entry("rdma", Bytes::from_static(b"rdma://10.0.0.1:6666"))
.unwrap();
builder
.add_entry("binary_data", Bytes::from_static(&[0x00, 0x01, 0x02, 0xFF]))
.unwrap();
let address = builder.build().unwrap();
// Verify WorkerAddress::get_entry() correctly decodes each entry
assert_eq!(
address.get_entry("tcp").unwrap(),
Some(Bytes::from_static(b"tcp://127.0.0.1:5555"))
);
assert_eq!(
address.get_entry("rdma").unwrap(),
Some(Bytes::from_static(b"rdma://10.0.0.1:6666"))
);
assert_eq!(
address.get_entry("binary_data").unwrap(),
Some(Bytes::from_static(&[0x00, 0x01, 0x02, 0xFF]))
);
assert_eq!(address.get_entry("nonexistent").unwrap(), None);
}
#[test]
fn test_builder_address_integration_available_transports() {
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry("tcp", Bytes::from_static(b"tcp://127.0.0.1:5555"))
.unwrap();
builder
.add_entry("rdma", Bytes::from_static(b"rdma://10.0.0.1:6666"))
.unwrap();
builder
.add_entry("grpc", Bytes::from_static(b"grpc://localhost:9000"))
.unwrap();
let address = builder.build().unwrap();
// Verify WorkerAddress::available_transports() returns all keys
let transports = address.available_transports().unwrap();
assert_eq!(transports.len(), 3);
assert!(transports.contains(&velo_common::TransportKey::from("tcp")));
assert!(transports.contains(&velo_common::TransportKey::from("rdma")));
assert!(transports.contains(&velo_common::TransportKey::from("grpc")));
}
#[test]
fn test_builder_address_integration_checksum_stability() {
// Build same address twice - checksums should match
let mut builder1 = WorkerAddressBuilder::new();
builder1
.add_entry("key", Bytes::from_static(b"value"))
.unwrap();
let address1 = builder1.build().unwrap();
let mut builder2 = WorkerAddressBuilder::new();
builder2
.add_entry("key", Bytes::from_static(b"value"))
.unwrap();
let address2 = builder2.build().unwrap();
// Same content should produce same checksum
assert_eq!(address1.checksum(), address2.checksum());
// Different content should produce different checksum
let mut builder3 = WorkerAddressBuilder::new();
builder3
.add_entry("key", Bytes::from_static(b"different"))
.unwrap();
let address3 = builder3.build().unwrap();
assert_ne!(address1.checksum(), address3.checksum());
}
#[test]
fn test_builder_address_integration_bytes_roundtrip() {
// Build an address
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry("endpoint", Bytes::from_static(b"test://value"))
.unwrap();
let address = builder.build().unwrap();
// Get raw bytes and create new address via from_encoded
let raw_bytes = address.to_bytes();
let address2 = WorkerAddress::from_encoded(raw_bytes);
// Both should be equal and decode the same
assert_eq!(address, address2);
assert_eq!(address.checksum(), address2.checksum());
assert_eq!(
address.get_entry("endpoint").unwrap(),
address2.get_entry("endpoint").unwrap()
);
}
#[test]
fn test_builder_address_integration_serde_roundtrip() {
// Build an address
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry("tcp", Bytes::from_static(b"tcp://127.0.0.1:5555"))
.unwrap();
let address = builder.build().unwrap();
// Serialize to JSON and back
let json = serde_json::to_string(&address).unwrap();
let deserialized: WorkerAddress = serde_json::from_str(&json).unwrap();
// Should be equal and decode correctly
assert_eq!(address, deserialized);
assert_eq!(
address.get_entry("tcp").unwrap(),
deserialized.get_entry("tcp").unwrap()
);
}
}
lib/velo-transports/src/lib.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
#![deny(missing_docs)]
//! Multi-transport active message routing framework.
//!
//! `velo-transports` abstracts TCP, HTTP, NATS, gRPC, and UCX behind a unified
//! [`Transport`] trait with zero-copy [`bytes::Bytes`], fire-and-forget error
//! callbacks, priority-based peer routing, and 3-phase graceful shutdown.
//!
//! # Architecture
//!
//! [`VeloBackend`] is the central orchestrator. It holds a set of transports,
//! each identified by a [`TransportKey`]. When a peer registers, the backend
//! selects a *primary* transport (highest-priority compatible transport) and
//! records any alternatives. Outbound messages are routed through the primary
//! transport by default, or through an explicit alternative.
//!
//! Inbound messages arrive via [`DataStreams`] — three independent channels
//! for messages, responses, and events.
//!
//! # Shutdown
//!
//! Graceful shutdown follows three phases:
//! 1. **Gate** — flip the draining flag; transports reject new inbound requests.
//! 2. **Drain** — wait for all in-flight requests to complete.
//! 3. **Teardown** — cancel listeners/writers and call `shutdown()` on each transport.
mod address;
pub mod tcp;
// #[cfg(feature = "ucx")]
// pub mod ucx;
// #[cfg(feature = "http")]
// pub mod http;
// #[cfg(feature = "nats")]
// pub mod nats;
// #[cfg(feature = "grpc")]
// pub mod grpc;
mod transport;
use std::{collections::HashMap, sync::Arc};
use dashmap::DashMap;
use parking_lot::Mutex;
// Public re-exports from velo-common
pub use velo_common::{
InstanceId, PeerInfo, TransportKey, WorkerAddress, WorkerAddressError, WorkerId,
};
// Internal builder for address construction
use address::WorkerAddressBuilder;
// Re-export transport types
pub use transport::{
DataStreams, HealthCheckError, InFlightGuard, MessageType, ShutdownPolicy, ShutdownState,
Transport, TransportAdapter, TransportError, TransportErrorHandler, make_channels,
};
/// Errors returned by [`VeloBackend`] operations.
#[derive(Debug, thiserror::Error)]
pub enum VeloBackendError {
/// No transport could accept the peer's address.
#[error("No compatible transports found")]
NoCompatibleTransports,
/// The target instance was never registered via [`VeloBackend::register_peer`].
#[error("Transport not found for instance: {0}")]
InstanceNotRegistered(InstanceId),
/// The worker ID is not in the fast-path cache.
#[error("Worker not found: {0}")]
WorkerNotRegistered(WorkerId),
/// The requested [`TransportKey`] does not match any loaded transport.
#[error("Transport not found: {0}")]
TransportNotFound(TransportKey),
/// The priority list does not match the set of available transports.
#[error("Invalid transport priority: {0}")]
InvalidTransportPriority(String),
}
/// Central orchestrator that aggregates multiple transports and routes messages
/// to peers via priority-based transport selection.
///
/// Each peer is registered with all compatible transports; the highest-priority
/// compatible transport becomes the *primary* for that peer. Worker IDs are
/// cached for fast-path routing without discovery lookups.
pub struct VeloBackend {
instance_id: InstanceId,
address: WorkerAddress,
priorities: Mutex<Vec<TransportKey>>,
transports: HashMap<TransportKey, Arc<dyn Transport>>,
primary_transport: DashMap<InstanceId, Arc<dyn Transport>>,
alternative_transports: DashMap<InstanceId, Vec<TransportKey>>,
workers: DashMap<WorkerId, InstanceId>,
shutdown_state: ShutdownState,
#[allow(dead_code)]
runtime: tokio::runtime::Handle,
}
impl VeloBackend {
/// Create a new backend from a list of transports.
///
/// Each transport is started (bound, listening) and its address is merged
/// into a composite [`WorkerAddress`]. Returns the backend and the
/// [`DataStreams`] receivers for inbound messages.
pub async fn new(
backend_transports: Vec<Arc<dyn Transport>>,
) -> anyhow::Result<(Self, DataStreams)> {
let instance_id = InstanceId::new_v4();
// build worker address
let mut priorities = Vec::new();
let mut builder = WorkerAddressBuilder::new();
let mut transports = HashMap::new();
let (adapter, data_streams) = transport::make_channels();
let shutdown_state = adapter.shutdown_state.clone();
let runtime = tokio::runtime::Handle::current();
for transport in backend_transports {
transport
.start(instance_id, adapter.clone(), runtime.clone())
.await?;
builder.merge(&transport.address())?;
priorities.push(transport.key());
transports.insert(transport.key(), transport);
}
let address = builder.build()?;
Ok((
Self {
instance_id,
address,
transports,
priorities: Mutex::new(priorities),
primary_transport: DashMap::new(),
alternative_transports: DashMap::new(),
workers: DashMap::new(),
shutdown_state,
runtime,
},
data_streams,
))
}
/// Returns this backend's unique instance identifier.
pub fn instance_id(&self) -> InstanceId {
self.instance_id
}
/// Returns a [`PeerInfo`] describing this backend (instance ID + composite address).
pub fn peer_info(&self) -> PeerInfo {
PeerInfo::new(self.instance_id, self.address.clone())
}
/// Returns `true` if the given instance has been registered via [`register_peer`](Self::register_peer).
pub fn is_registered(&self, instance_id: InstanceId) -> bool {
self.primary_transport.contains_key(&instance_id)
}
/// Fast-path lookup of worker_id -> instance_id from cache.
///
/// Returns `WorkerNotRegistered` if the worker is not in the cache.
/// Higher layers (Velo, VeloEvents, ActiveMessageClient) should handle
/// discovery fallback when this returns an error.
///
/// # Example
/// ```ignore
/// match backend.try_translate_worker_id(worker_id) {
/// Ok(instance_id) => { /* fast path: send immediately */ }
/// Err(VeloBackendError::WorkerNotRegistered(_)) => {
/// /* slow path: query discovery, then register_peer() */
/// }
/// }
/// ```
pub fn try_translate_worker_id(
&self,
worker_id: WorkerId,
) -> Result<InstanceId, VeloBackendError> {
self.workers
.get(&worker_id)
.map(|entry| *entry)
.ok_or(VeloBackendError::WorkerNotRegistered(worker_id))
}
/// Deprecated: Use `try_translate_worker_id()` for explicit fast-path semantics.
#[deprecated(since = "0.7.0", note = "Use try_translate_worker_id() instead")]
pub fn translate_worker_id(&self, worker_id: WorkerId) -> Result<InstanceId, VeloBackendError> {
self.try_translate_worker_id(worker_id)
}
/// Check if an instance_id is registered.
pub fn has_instance(&self, instance_id: InstanceId) -> bool {
self.primary_transport.contains_key(&instance_id)
}
/// Send a message to a registered peer via its primary transport.
///
/// Returns [`VeloBackendError::InstanceNotRegistered`] if the peer has not
/// been registered with [`register_peer`](Self::register_peer).
pub fn send_message(
&self,
target: InstanceId,
header: Vec<u8>,
payload: Vec<u8>,
message_type: MessageType,
on_error: Arc<dyn TransportErrorHandler>,
) -> anyhow::Result<()> {
let transport = self
.primary_transport
.get(&target)
.ok_or(VeloBackendError::InstanceNotRegistered(target))?;
transport.send_message(target, header, payload, message_type, on_error);
Ok(())
}
/// Send a message to a registered peer via a specific transport.
///
/// If `transport_key` matches the peer's primary transport, the message is
/// sent directly. Otherwise, the alternative transports are searched.
/// Returns [`VeloBackendError::NoCompatibleTransports`] if the requested
/// transport is not available for this peer.
pub fn send_message_with_transport(
&self,
target: InstanceId,
header: Vec<u8>,
payload: Vec<u8>,
message_type: MessageType,
on_error: Arc<dyn TransportErrorHandler>,
transport_key: TransportKey,
) -> anyhow::Result<()> {
let transport = self
.primary_transport
.get(&target)
.ok_or(VeloBackendError::InstanceNotRegistered(target))?;
if transport.value().key() == transport_key {
transport.send_message(target, header, payload, message_type, on_error);
return Ok(());
} else {
// if we got here, we can unwrap because there is an entry in the alternative_transports map
let alternative_transports = self
.alternative_transports
.get(&target)
.ok_or(VeloBackendError::InstanceNotRegistered(target))?;
for alternative_transport in alternative_transports.iter() {
if *alternative_transport == transport_key
&& let Some(transport) = self.transports.get(alternative_transport)
{
transport.send_message(target, header, payload, message_type, on_error);
return Ok(());
}
}
}
Err(VeloBackendError::NoCompatibleTransports)?
}
/// Send message to a worker (fast-path only).
///
/// This method uses `try_translate_worker_id()` for fast-path lookup.
/// Returns `WorkerNotRegistered` error if the worker is not in the cache.
///
/// For automatic discovery, use the two-phase pattern:
/// ```ignore
/// match backend.send_message_to_worker(...) {
/// Ok(()) => { /* success */ }
/// Err(e) if matches_worker_not_registered(&e) => {
/// tokio::spawn(async move {
/// let instance_id = backend.resolve_and_register_worker(worker_id).await?;
/// backend.send_message(instance_id, ...)?;
/// });
/// }
/// }
/// ```
pub fn send_message_to_worker(
&self,
worker_id: WorkerId,
header: Vec<u8>,
payload: Vec<u8>,
message_type: MessageType,
on_error: Arc<dyn TransportErrorHandler>,
) -> anyhow::Result<()> {
let instance_id = self.try_translate_worker_id(worker_id)?;
self.send_message(instance_id, header, payload, message_type, on_error)
}
/// Register a remote peer with all compatible transports.
///
/// The highest-priority compatible transport becomes the peer's *primary*.
/// Returns [`VeloBackendError::NoCompatibleTransports`] if no transport
/// can accept the peer's address.
pub fn register_peer(&self, peer: PeerInfo) -> Result<(), VeloBackendError> {
// try to register the peer with each transport
// we must have at least one compatible transport; otherwise, return an error
let instance_id = peer.instance_id();
let mut compatible_transports = Vec::new();
for (key, transport) in self.transports.iter() {
if transport.register(peer.clone()).is_ok() {
compatible_transports.push(key.clone());
}
}
if compatible_transports.is_empty() {
return Err(VeloBackendError::NoCompatibleTransports);
}
// sort against the preferred transports
let sorted_transports = self
.priorities
.lock()
.iter()
.filter(|key| compatible_transports.contains(key))
.cloned()
.collect::<Vec<TransportKey>>();
assert!(
!sorted_transports.is_empty(),
"failed to properly sort compatible transports"
);
let primary_transport_key = sorted_transports[0].clone();
let alternative_transport_keys = sorted_transports[1..].to_vec();
let primary_transport = self.transports.get(&primary_transport_key).unwrap();
self.primary_transport
.insert(instance_id, primary_transport.clone());
self.alternative_transports
.insert(instance_id, alternative_transport_keys);
self.workers.insert(instance_id.worker_id(), instance_id);
Ok(())
}
/// Get the available transports.
pub fn available_transports(&self) -> Vec<TransportKey> {
self.transports.keys().cloned().collect()
}
/// Set the priority of the transports.
///
/// The list of [`TransportKey`]s must be an order set of the available transports.
pub fn set_transport_priority(
&self,
priorities: Vec<TransportKey>,
) -> Result<(), VeloBackendError> {
let required_transports = self.available_transports();
if required_transports.len() != priorities.len() {
return Err(VeloBackendError::InvalidTransportPriority(format!(
"Required transports: {:?}, provided priorities: {:?}",
required_transports, priorities
)));
}
for priority in &priorities {
if !required_transports.contains(priority) {
return Err(VeloBackendError::InvalidTransportPriority(format!(
"Priority transport not found: {:?}",
priority
)));
}
}
let mut guard = self.priorities.lock();
*guard = priorities;
Ok(())
}
/// Get the shared shutdown state.
pub fn shutdown_state(&self) -> &ShutdownState {
&self.shutdown_state
}
/// Perform a graceful 3-phase shutdown.
///
/// 1. **Gate**: Flip the draining flag and notify each transport via `begin_drain()`.
/// 2. **Drain**: Wait for all in-flight requests to complete (per `policy`).
/// 3. **Teardown**: Cancel the teardown token and call `shutdown()` on each transport.
pub async fn graceful_shutdown(&self, policy: ShutdownPolicy) {
// Phase 1: Gate
self.shutdown_state.begin_drain();
for transport in self.transports.values() {
transport.begin_drain();
}
// Phase 2: Drain
match policy {
ShutdownPolicy::WaitForever => {
self.shutdown_state.wait_for_drain().await;
}
ShutdownPolicy::Timeout(duration) => {
let _ = tokio::time::timeout(duration, self.shutdown_state.wait_for_drain()).await;
}
}
// Phase 3: Teardown
self.shutdown_state.teardown_token().cancel();
for transport in self.transports.values() {
transport.shutdown();
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use bytes::Bytes;
use futures::future::BoxFuture;
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::time::Duration;
/// Mock transport for testing VeloBackend logic without real networking.
struct MockTransport {
key: TransportKey,
address: WorkerAddress,
accept_register: bool,
started: AtomicBool,
drained: AtomicBool,
shut_down: AtomicBool,
send_count: AtomicUsize,
}
impl MockTransport {
fn new(key: &str, accept_register: bool) -> Arc<Self> {
let mut builder = WorkerAddressBuilder::new();
builder
.add_entry(key, format!("mock://{}", key).into_bytes())
.unwrap();
let address = builder.build().unwrap();
Arc::new(Self {
key: TransportKey::from(key),
address,
accept_register,
started: AtomicBool::new(false),
drained: AtomicBool::new(false),
shut_down: AtomicBool::new(false),
send_count: AtomicUsize::new(0),
})
}
}
impl Transport for MockTransport {
fn key(&self) -> TransportKey {
self.key.clone()
}
fn address(&self) -> WorkerAddress {
self.address.clone()
}
fn register(&self, _peer_info: PeerInfo) -> Result<(), TransportError> {
if self.accept_register {
Ok(())
} else {
Err(TransportError::NoEndpoint)
}
}
fn send_message(
&self,
_instance_id: InstanceId,
_header: Vec<u8>,
_payload: Vec<u8>,
_message_type: MessageType,
_on_error: Arc<dyn TransportErrorHandler>,
) {
self.send_count.fetch_add(1, Ordering::Relaxed);
}
fn start(
&self,
_instance_id: InstanceId,
_channels: TransportAdapter,
_rt: tokio::runtime::Handle,
) -> BoxFuture<'_, anyhow::Result<()>> {
self.started.store(true, Ordering::Relaxed);
Box::pin(async { Ok(()) })
}
fn shutdown(&self) {
self.shut_down.store(true, Ordering::Relaxed);
}
fn begin_drain(&self) {
self.drained.store(true, Ordering::Relaxed);
}
fn check_health(
&self,
_instance_id: InstanceId,
_timeout: Duration,
) -> std::pin::Pin<
Box<
dyn std::future::Future<Output = Result<(), transport::HealthCheckError>>
+ Send
+ '_,
>,
> {
Box::pin(async { Ok(()) })
}
}
struct NoopErrorHandler;
impl TransportErrorHandler for NoopErrorHandler {
fn on_error(&self, _header: Bytes, _payload: Bytes, _error: String) {}
}
/// Helper: build a PeerInfo with entries for specified transport keys.
fn make_peer_info(keys: &[&str]) -> PeerInfo {
let instance_id = InstanceId::new_v4();
let mut builder = WorkerAddressBuilder::new();
for key in keys {
builder
.add_entry(*key, format!("mock://{}", key).into_bytes())
.unwrap();
}
let address = builder.build().unwrap();
PeerInfo::new(instance_id, address)
}
#[tokio::test]
async fn test_new_single_transport() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t.clone() as Arc<dyn Transport>])
.await
.unwrap();
assert!(t.started.load(Ordering::Relaxed));
// instance_id should be a valid v4 UUID (non-zero)
assert!(!backend.instance_id().as_bytes().iter().all(|&b| b == 0));
assert_eq!(backend.available_transports().len(), 1);
}
#[tokio::test]
async fn test_new_multiple_transports() {
let t1 = MockTransport::new("tcp", true);
let t2 = MockTransport::new("http", true);
let (backend, _streams) = VeloBackend::new(vec![
t1.clone() as Arc<dyn Transport>,
t2.clone() as Arc<dyn Transport>,
])
.await
.unwrap();
assert!(t1.started.load(Ordering::Relaxed));
assert!(t2.started.load(Ordering::Relaxed));
assert_eq!(backend.available_transports().len(), 2);
}
#[tokio::test]
async fn test_register_peer_selects_primary_by_priority() {
let t1 = MockTransport::new("tcp", true);
let t2 = MockTransport::new("http", true);
let (backend, _streams) = VeloBackend::new(vec![
t1.clone() as Arc<dyn Transport>,
t2.clone() as Arc<dyn Transport>,
])
.await
.unwrap();
let peer = make_peer_info(&["tcp", "http"]);
let peer_id = peer.instance_id();
backend.register_peer(peer).unwrap();
assert!(backend.is_registered(peer_id));
// Primary should be "tcp" (first in priority)
let primary = backend.primary_transport.get(&peer_id).unwrap();
assert_eq!(primary.value().key(), TransportKey::from("tcp"));
}
#[tokio::test]
async fn test_register_peer_no_compatible_transports() {
// Transport rejects all registrations
let t = MockTransport::new("tcp", false);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let peer = make_peer_info(&["tcp"]);
let result = backend.register_peer(peer);
assert!(matches!(
result,
Err(VeloBackendError::NoCompatibleTransports)
));
}
#[tokio::test]
async fn test_register_peer_stores_worker_mapping() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let peer = make_peer_info(&["tcp"]);
let peer_id = peer.instance_id();
let worker_id = peer_id.worker_id();
backend.register_peer(peer).unwrap();
let resolved = backend.try_translate_worker_id(worker_id).unwrap();
assert_eq!(resolved, peer_id);
}
#[tokio::test]
async fn test_send_message_routes_to_primary() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t.clone() as Arc<dyn Transport>])
.await
.unwrap();
let peer = make_peer_info(&["tcp"]);
let peer_id = peer.instance_id();
backend.register_peer(peer).unwrap();
backend
.send_message(
peer_id,
vec![1],
vec![2],
MessageType::Message,
Arc::new(NoopErrorHandler),
)
.unwrap();
assert_eq!(t.send_count.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn test_send_message_unregistered_peer() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let result = backend.send_message(
InstanceId::new_v4(),
vec![],
vec![],
MessageType::Message,
Arc::new(NoopErrorHandler),
);
assert!(result.is_err());
}
#[tokio::test]
async fn test_send_message_with_transport_primary_match() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t.clone() as Arc<dyn Transport>])
.await
.unwrap();
let peer = make_peer_info(&["tcp"]);
let peer_id = peer.instance_id();
backend.register_peer(peer).unwrap();
backend
.send_message_with_transport(
peer_id,
vec![1],
vec![2],
MessageType::Message,
Arc::new(NoopErrorHandler),
TransportKey::from("tcp"),
)
.unwrap();
assert_eq!(t.send_count.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn test_send_message_with_transport_alternative() {
let t1 = MockTransport::new("tcp", true);
let t2 = MockTransport::new("http", true);
let (backend, _streams) = VeloBackend::new(vec![
t1.clone() as Arc<dyn Transport>,
t2.clone() as Arc<dyn Transport>,
])
.await
.unwrap();
let peer = make_peer_info(&["tcp", "http"]);
let peer_id = peer.instance_id();
backend.register_peer(peer).unwrap();
// Send via "http" (the alternative transport)
backend
.send_message_with_transport(
peer_id,
vec![1],
vec![2],
MessageType::Message,
Arc::new(NoopErrorHandler),
TransportKey::from("http"),
)
.unwrap();
assert_eq!(t2.send_count.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn test_send_message_with_transport_not_found() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let peer = make_peer_info(&["tcp"]);
let peer_id = peer.instance_id();
backend.register_peer(peer).unwrap();
let result = backend.send_message_with_transport(
peer_id,
vec![],
vec![],
MessageType::Message,
Arc::new(NoopErrorHandler),
TransportKey::from("grpc"),
);
assert!(result.is_err());
}
#[tokio::test]
async fn test_try_translate_worker_id_not_found() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let result = backend.try_translate_worker_id(InstanceId::new_v4().worker_id());
assert!(matches!(
result,
Err(VeloBackendError::WorkerNotRegistered(_))
));
}
#[tokio::test]
async fn test_set_transport_priority_valid() {
let t1 = MockTransport::new("tcp", true);
let t2 = MockTransport::new("http", true);
let (backend, _streams) =
VeloBackend::new(vec![t1 as Arc<dyn Transport>, t2 as Arc<dyn Transport>])
.await
.unwrap();
// Reverse the priority
backend
.set_transport_priority(vec![TransportKey::from("http"), TransportKey::from("tcp")])
.unwrap();
}
#[tokio::test]
async fn test_set_transport_priority_wrong_length() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let result = backend
.set_transport_priority(vec![TransportKey::from("tcp"), TransportKey::from("http")]);
assert!(matches!(
result,
Err(VeloBackendError::InvalidTransportPriority(_))
));
}
#[tokio::test]
async fn test_set_transport_priority_unknown_key() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let result = backend.set_transport_priority(vec![TransportKey::from("unknown")]);
assert!(matches!(
result,
Err(VeloBackendError::InvalidTransportPriority(_))
));
}
#[tokio::test]
async fn test_graceful_shutdown_calls_all_transports() {
let t1 = MockTransport::new("tcp", true);
let t2 = MockTransport::new("http", true);
let (backend, _streams) = VeloBackend::new(vec![
t1.clone() as Arc<dyn Transport>,
t2.clone() as Arc<dyn Transport>,
])
.await
.unwrap();
backend
.graceful_shutdown(ShutdownPolicy::Timeout(Duration::from_millis(100)))
.await;
assert!(t1.drained.load(Ordering::Relaxed));
assert!(t2.drained.load(Ordering::Relaxed));
assert!(t1.shut_down.load(Ordering::Relaxed));
assert!(t2.shut_down.load(Ordering::Relaxed));
assert!(backend.shutdown_state().is_draining());
assert!(backend.shutdown_state().teardown_token().is_cancelled());
}
#[tokio::test]
async fn test_peer_info_roundtrip() {
let t = MockTransport::new("tcp", true);
let (backend, _streams) = VeloBackend::new(vec![t as Arc<dyn Transport>])
.await
.unwrap();
let info = backend.peer_info();
assert_eq!(info.instance_id(), backend.instance_id());
}
}
lib/velo-transports/src/tcp/framing.rs 0 → 100644
View file @ 63d7c01c
// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//! Zero-copy TCP framing codec for ActiveMessage transport
//!
//! Wire format (7-15 bytes overhead):
//! ```text
//! [u16 BE: schema_version][u8: frame_type][u32 BE: header_len][u32 BE: payload_len][header bytes][payload bytes]
//! ```
//!
//! The codec uses `BytesMut` for receiving and `Bytes` for output, enabling
//! zero-copy buffer slicing where header and payload share the underlying buffer.
use bytes::{Buf, Bytes, BytesMut};
use std::io;
use std::io::Write;
use tokio::io::{AsyncWrite, AsyncWriteExt};
use tokio_util::codec::Decoder;
use crate::MessageType;
/// Current schema version
const SCHEMA_VERSION_V1: u16 = 1;
/// Maximum frame size (16 MB)
const MAX_FRAME_SIZE: u32 = 16 * 1024 * 1024;
/// Minimum frame header size (version + type + 2 lengths)
const MIN_HEADER_SIZE: usize = 2 + 1 + 4 + 4; // 11 bytes
/// Zero-copy frame decoder for TCP transport
///
/// This decoder maintains state across multiple calls to support partial
/// frame reception. It decodes frames into (MessageType, header: Bytes, payload: Bytes)
/// where header and payload are zero-copy slices of the receive buffer.
#[derive(Debug, Clone)]
pub struct TcpFrameCodec {
state: DecodeState,
}
#[derive(Debug, Clone, Copy)]
enum DecodeState {
/// Waiting for frame header (version + type + lengths)
AwaitingHeader,
/// Waiting for frame data (header + payload), with known lengths
AwaitingData {
frame_type: MessageType,
header_len: u32,
payload_len: u32,
},
}
impl TcpFrameCodec {
/// Create a new frame codec
pub fn new() -> Self {
Self {
state: DecodeState::AwaitingHeader,
}
}
/// Build the frame preamble (metadata header)
///
/// Returns a fixed-size preamble containing version, message type, and lengths.
#[inline]
pub fn build_preamble(
msg_type: MessageType,
header_len: u32,
payload_len: u32,
) -> io::Result<[u8; MIN_HEADER_SIZE]> {
// Validate lengths before building preamble
Self::validate_lengths(header_len, payload_len)?;
let mut preamble = [0u8; MIN_HEADER_SIZE];
// Layout:
// [0..2) = version
// [2] = msg_type
// [3..7) = header_len
// [7..11)= payload_len (total 11 bytes)
preamble[0..2].copy_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
preamble[2] = msg_type.as_u8();
preamble[3..7].copy_from_slice(&header_len.to_be_bytes());
preamble[7..11].copy_from_slice(&payload_len.to_be_bytes());
Ok(preamble)
}
/// Parse message type from a preamble
///
/// Validates the schema version and extracts the message type from the preamble.
#[inline]
pub fn parse_message_type_from_preamble(preamble: &[u8]) -> io::Result<MessageType> {
if preamble.len() < MIN_HEADER_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"Preamble too short",
));
}
// Validate schema version
let schema_version = u16::from_be_bytes([preamble[0], preamble[1]]);
if schema_version != SCHEMA_VERSION_V1 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"Unsupported schema version: {} (expected {})",
schema_version, SCHEMA_VERSION_V1
),
));
}
// Extract and validate message type
MessageType::from_u8(preamble[2]).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("Invalid message type: {}", preamble[2]),
)
})
}
/// Encode and write a frame asynchronously
///
/// Uses `write_all()` for each segment to handle partial writes correctly.
/// TCP `write_vectored()` doesn't guarantee writing all bytes in one call —
/// for payloads exceeding the kernel send buffer (~128KB), it returns a short
/// write count. Using `write_all()` per segment ensures correctness for all sizes.
#[inline]
pub async fn encode_frame<W: AsyncWrite + Unpin>(
writer: &mut W,
msg_type: MessageType,
header: &[u8],
payload: &[u8],
) -> tokio::io::Result<()> {
let preamble = Self::build_preamble(msg_type, header.len() as u32, payload.len() as u32)?;
writer.write_all(&preamble).await?;
writer.write_all(header).await?;
writer.write_all(payload).await?;
Ok(())
}
/// Encode and write a frame synchronously
///
/// Uses `write_all()` for each segment to handle partial writes correctly.
#[inline]
pub fn encode_frame_sync<W: Write>(
writer: &mut W,
msg_type: MessageType,
header: &[u8],
payload: &[u8],
) -> std::io::Result<()> {
let preamble = Self::build_preamble(msg_type, header.len() as u32, payload.len() as u32)?;
writer.write_all(&preamble)?;
writer.write_all(header)?;
writer.write_all(payload)?;
Ok(())
}
/// Validate that lengths are reasonable
fn validate_lengths(header_len: u32, payload_len: u32) -> io::Result<()> {
let total_len = header_len
.checked_add(payload_len)
.ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Frame size overflow"))?;
if total_len > MAX_FRAME_SIZE {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"Frame size {} exceeds maximum {}",
total_len, MAX_FRAME_SIZE
),
));
}
Ok(())
}
}
impl Default for TcpFrameCodec {
fn default() -> Self {
Self::new()
}
}
impl Decoder for TcpFrameCodec {
type Item = (MessageType, Bytes, Bytes);
type Error = io::Error;
fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
loop {
match self.state {
DecodeState::AwaitingHeader => {
// Need at least MIN_HEADER_SIZE bytes
if src.len() < MIN_HEADER_SIZE {
return Ok(None);
}
// Parse header without consuming bytes yet
let schema_version = u16::from_be_bytes([src[0], src[1]]);
let frame_type_byte = src[2];
let header_len = u32::from_be_bytes([src[3], src[4], src[5], src[6]]);
let payload_len = u32::from_be_bytes([src[7], src[8], src[9], src[10]]);
// Validate schema version
if schema_version != SCHEMA_VERSION_V1 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"Unsupported schema version: {} (expected {})",
schema_version, SCHEMA_VERSION_V1
),
));
}
// Parse frame type
let frame_type = MessageType::from_u8(frame_type_byte).ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidData,
format!("Invalid frame type: {}", frame_type_byte),
)
})?;
// Validate lengths before allocating/waiting
Self::validate_lengths(header_len, payload_len)?;
// Advance buffer past header
src.advance(MIN_HEADER_SIZE);
// Transition to data state
self.state = DecodeState::AwaitingData {
frame_type,
header_len,
payload_len,
};
}
DecodeState::AwaitingData {
frame_type,
header_len,
payload_len,
..
} => {
let total_data_len = (header_len + payload_len) as usize;
// Wait for full data
if src.len() < total_data_len {
return Ok(None);
}
// Zero-copy: split buffer into header and payload slices
let header = src.split_to(header_len as usize).freeze();
let payload = src.split_to(payload_len as usize).freeze();
// Reset state for next frame
self.state = DecodeState::AwaitingHeader;
return Ok(Some((frame_type, header, payload)));
}
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Test helper to encode a frame into a Vec<u8> for verification (async)
async fn encode_frame_to_bytes(
msg_type: MessageType,
header: &[u8],
payload: &[u8],
) -> io::Result<Vec<u8>> {
let mut buf = Vec::new();
TcpFrameCodec::encode_frame(&mut buf, msg_type, header, payload).await?;
Ok(buf)
}
/// Test helper to encode a frame into a Vec<u8> for verification (sync)
fn encode_frame_to_bytes_sync(
msg_type: MessageType,
header: &[u8],
payload: &[u8],
) -> io::Result<Vec<u8>> {
let mut buf = Vec::new();
TcpFrameCodec::encode_frame_sync(&mut buf, msg_type, header, payload)?;
Ok(buf)
}
/// Helper to create raw frames with arbitrary parameters for negative testing.
///
/// This function bypasses normal validation and encoding logic to create
/// intentionally invalid frames (wrong schema version, oversized frames, etc.)
/// for testing error handling paths. Use `encode_frame_to_bytes()` for
/// testing valid frame construction.
fn create_unsafe_frame(
schema_version: u16,
frame_type: MessageType,
header: &[u8],
payload: &[u8],
) -> BytesMut {
let mut buf = BytesMut::new();
buf.extend_from_slice(&schema_version.to_be_bytes());
buf.extend_from_slice(&[frame_type.as_u8()]);
buf.extend_from_slice(&(header.len() as u32).to_be_bytes());
buf.extend_from_slice(&(payload.len() as u32).to_be_bytes());
buf.extend_from_slice(header);
buf.extend_from_slice(payload);
buf
}
#[test]
fn test_decode_message_frame() {
let mut codec = TcpFrameCodec::new();
let header = b"test-header";
let payload = b"test-payload-data";
let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
let mut buf = BytesMut::from(&framed[..]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (msg_type, decoded_header, decoded_payload) = result.unwrap();
assert_eq!(msg_type, MessageType::Message);
assert_eq!(decoded_header, Bytes::from(header.as_ref()));
assert_eq!(decoded_payload, Bytes::from(payload.as_ref()));
}
#[test]
fn test_decode_all_frame_types() {
let frame_types = [
MessageType::Message,
MessageType::Response,
MessageType::Ack,
MessageType::Event,
];
for frame_type in &frame_types {
let mut codec = TcpFrameCodec::new();
let header = b"header";
let payload = b"payload";
let framed = encode_frame_to_bytes_sync(*frame_type, header, payload).unwrap();
let mut buf = BytesMut::from(&framed[..]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (decoded_type, _, _) = result.unwrap();
assert_eq!(decoded_type, *frame_type);
}
}
#[test]
fn test_decode_empty_payload() {
let mut codec = TcpFrameCodec::new();
let header = b"ack-header";
let payload = b"";
let framed = encode_frame_to_bytes_sync(MessageType::Ack, header, payload).unwrap();
let mut buf = BytesMut::from(&framed[..]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (msg_type, decoded_header, decoded_payload) = result.unwrap();
assert_eq!(msg_type, MessageType::Ack);
assert_eq!(&decoded_header[..], header);
assert_eq!(decoded_payload.len(), 0);
}
#[test]
fn test_decode_partial_frame() {
let mut codec = TcpFrameCodec::new();
let header = b"test-header";
let payload = b"test-payload";
let full_frame = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
// Send only first 5 bytes (partial header)
let mut buf = BytesMut::from(&full_frame[..5]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_none()); // Not enough data
// Send rest of header
buf.extend_from_slice(&full_frame[5..MIN_HEADER_SIZE]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_none()); // Header parsed, but data not yet available
// Send complete data
buf.extend_from_slice(&full_frame[MIN_HEADER_SIZE..]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (msg_type, decoded_header, decoded_payload) = result.unwrap();
assert_eq!(msg_type, MessageType::Message);
assert_eq!(&decoded_header[..], header);
assert_eq!(&decoded_payload[..], payload);
}
#[test]
fn test_decode_invalid_schema_version() {
let mut codec = TcpFrameCodec::new();
let header = b"header";
let payload = b"payload";
let mut buf = create_unsafe_frame(999, MessageType::Message, header, payload);
let result = codec.decode(&mut buf);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Unsupported schema version")
);
}
#[test]
fn test_decode_invalid_frame_type() {
let mut codec = TcpFrameCodec::new();
let mut buf = BytesMut::new();
// Create frame with invalid type byte (255)
buf.extend_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
buf.extend_from_slice(&[255u8]); // Invalid frame type
buf.extend_from_slice(&10u32.to_be_bytes()); // header len
buf.extend_from_slice(&10u32.to_be_bytes()); // payload len
let result = codec.decode(&mut buf);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.to_string()
.contains("Invalid frame type")
);
}
#[test]
fn test_decode_frame_too_large() {
let mut codec = TcpFrameCodec::new();
let mut buf = BytesMut::new();
// Create frame that exceeds MAX_FRAME_SIZE
buf.extend_from_slice(&SCHEMA_VERSION_V1.to_be_bytes());
buf.extend_from_slice(&[MessageType::Message.as_u8()]);
buf.extend_from_slice(&(MAX_FRAME_SIZE / 2 + 1).to_be_bytes());
buf.extend_from_slice(&(MAX_FRAME_SIZE / 2 + 1).to_be_bytes());
let result = codec.decode(&mut buf);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("exceeds maximum"));
}
#[test]
fn test_decode_multiple_frames() {
let mut codec = TcpFrameCodec::new();
let mut buf = BytesMut::new();
// Add two frames to buffer
let frame1 =
encode_frame_to_bytes_sync(MessageType::Message, b"header1", b"payload1").unwrap();
let frame2 =
encode_frame_to_bytes_sync(MessageType::Response, b"header2", b"payload2").unwrap();
buf.extend_from_slice(&frame1);
buf.extend_from_slice(&frame2);
// Decode first frame
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (msg_type, header, payload) = result.unwrap();
assert_eq!(msg_type, MessageType::Message);
assert_eq!(&header[..], b"header1");
assert_eq!(&payload[..], b"payload1");
// Decode second frame
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (msg_type, header, payload) = result.unwrap();
assert_eq!(msg_type, MessageType::Response);
assert_eq!(&header[..], b"header2");
assert_eq!(&payload[..], b"payload2");
// No more frames
assert!(buf.is_empty());
}
#[test]
fn test_zero_copy_bytes_share_buffer() {
let mut codec = TcpFrameCodec::new();
let header = b"shared-header";
let payload = b"shared-payload";
let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
let mut buf = BytesMut::from(&framed[..]);
let result = codec.decode(&mut buf).unwrap().unwrap();
let (_, decoded_header, decoded_payload) = result;
// Verify the slices contain correct data
assert_eq!(&decoded_header[..], header);
assert_eq!(&decoded_payload[..], payload);
// Clone should be cheap (just RC increment)
let header_clone = decoded_header.clone();
let payload_clone = decoded_payload.clone();
assert_eq!(decoded_header, header_clone);
assert_eq!(decoded_payload, payload_clone);
}
#[test]
fn test_encode_frame() {
let header = b"test-header";
let payload = b"test-payload";
let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
// Verify frame structure
assert_eq!(framed.len(), MIN_HEADER_SIZE + header.len() + payload.len());
// Verify header fields
assert_eq!(
u16::from_be_bytes([framed[0], framed[1]]),
SCHEMA_VERSION_V1
);
assert_eq!(framed[2], MessageType::Message.as_u8());
assert_eq!(
u32::from_be_bytes([framed[3], framed[4], framed[5], framed[6]]),
header.len() as u32
);
assert_eq!(
u32::from_be_bytes([framed[7], framed[8], framed[9], framed[10]]),
payload.len() as u32
);
// Verify data
assert_eq!(
&framed[MIN_HEADER_SIZE..MIN_HEADER_SIZE + header.len()],
header
);
assert_eq!(&framed[MIN_HEADER_SIZE + header.len()..], payload);
}
#[test]
fn test_encode_all_message_types() {
let header = b"header";
let payload = b"payload";
for msg_type in &[
MessageType::Message,
MessageType::Response,
MessageType::Ack,
MessageType::Event,
] {
let framed = encode_frame_to_bytes_sync(*msg_type, header, payload).unwrap();
assert_eq!(framed[2], msg_type.as_u8());
}
}
#[test]
fn test_encode_empty_payload() {
let header = b"ack-header";
let payload = b"";
let framed = encode_frame_to_bytes_sync(MessageType::Ack, header, payload).unwrap();
assert_eq!(framed.len(), MIN_HEADER_SIZE + header.len());
assert_eq!(
u32::from_be_bytes([framed[7], framed[8], framed[9], framed[10]]),
0
);
}
#[test]
fn test_encode_frame_too_large() {
let header = vec![0u8; (MAX_FRAME_SIZE / 2 + 1) as usize];
let payload = vec![0u8; (MAX_FRAME_SIZE / 2 + 1) as usize];
let result = encode_frame_to_bytes_sync(MessageType::Message, &header, &payload);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("exceeds maximum"));
}
#[test]
fn test_round_trip_encode_decode() {
let mut codec = TcpFrameCodec::new();
let header = b"round-trip-header";
let payload = b"round-trip-payload-data";
// Encode
let framed = encode_frame_to_bytes_sync(MessageType::Response, header, payload).unwrap();
// Decode
let mut buf = BytesMut::from(&framed[..]);
let result = codec.decode(&mut buf).unwrap();
assert!(result.is_some());
let (msg_type, decoded_header, decoded_payload) = result.unwrap();
assert_eq!(msg_type, MessageType::Response);
assert_eq!(&decoded_header[..], header);
assert_eq!(&decoded_payload[..], payload);
}
#[test]
fn test_round_trip_all_types() {
let types = [
MessageType::Message,
MessageType::Response,
MessageType::Ack,
MessageType::Event,
];
for msg_type in &types {
let mut codec = TcpFrameCodec::new();
let header = b"header";
let payload = b"payload";
let framed = encode_frame_to_bytes_sync(*msg_type, header, payload).unwrap();
let mut buf = BytesMut::from(&framed[..]);
let result = codec.decode(&mut buf).unwrap().unwrap();
assert_eq!(result.0, *msg_type);
assert_eq!(&result.1[..], header);
assert_eq!(&result.2[..], payload);
}
}
#[test]
fn test_encode_frame_sync() {
let header = b"sync-header";
let payload = b"sync-payload";
let framed = encode_frame_to_bytes_sync(MessageType::Message, header, payload).unwrap();
// Verify frame structure
assert_eq!(framed.len(), MIN_HEADER_SIZE + header.len() + payload.len());
// Verify preamble fields
assert_eq!(
u16::from_be_bytes([framed[0], framed[1]]),
SCHEMA_VERSION_V1
);
assert_eq!(framed[2], MessageType::Message.as_u8());
assert_eq!(
u32::from_be_bytes([framed[3], framed[4], framed[5], framed[6]]),
header.len() as u32
);
assert_eq!(
u32::from_be_bytes([framed[7], framed[8], framed[9], framed[10]]),
payload.len() as u32
);
// Verify data
assert_eq!(
&framed[MIN_HEADER_SIZE..MIN_HEADER_SIZE + header.len()],
header
);
assert_eq!(&framed[MIN_HEADER_SIZE + header.len()..], payload);
}
#[test]
fn test_sync_async_produce_same_output() {
let header = b"test-header";
let payload = b"test-payload";
// Encode with sync version
let sync_framed =
encode_frame_to_bytes_sync(MessageType::Response, header, payload).unwrap();
// Encode with async version (using tokio runtime)
let async_framed = tokio::runtime::Runtime::new()
.unwrap()
.block_on(encode_frame_to_bytes(
MessageType::Response,
header,
payload,
))
.unwrap();
// Both should produce identical output
assert_eq!(sync_framed, async_framed);
}
}
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment