velo_ext/id/
transport.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Transport key type for type-safe transport identification.
5
6use serde::{Deserialize, Serialize};
7use std::fmt;
8use std::sync::Arc;
9
10/// A type-safe wrapper around transport keys for WorkerAddress.
11///
12/// This provides a zero-cost abstraction over `Arc<str>` with type safety
13/// to prevent accidentally mixing transport keys with other string types.
14///
15/// # Examples
16///
17/// ```
18/// use velo_ext::TransportKey;
19///
20/// let key = TransportKey::new("tcp");
21/// assert_eq!(key.as_str(), "tcp");
22///
23/// // Ergonomic conversions
24/// let key2: TransportKey = "rdma".into();
25/// let key3 = TransportKey::from("udp");
26///
27/// // Use in collections
28/// use std::collections::HashMap;
29/// let mut transports = HashMap::new();
30/// transports.insert(TransportKey::from("tcp"), "tcp://127.0.0.1:5555");
31/// ```
32#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
33#[serde(transparent)]
34pub struct TransportKey(Arc<str>);
35
36impl TransportKey {
37    /// Create a new TransportKey from any type that can be converted into `Arc<str>`.
38    pub fn new(key: impl Into<Arc<str>>) -> Self {
39        Self(key.into())
40    }
41
42    /// Get the key as a string slice.
43    pub fn as_str(&self) -> &str {
44        &self.0
45    }
46}
47
48// Deref to str for ergonomic usage
49impl std::ops::Deref for TransportKey {
50    type Target = str;
51
52    fn deref(&self) -> &Self::Target {
53        &self.0
54    }
55}
56
57// AsRef for flexible parameter types
58impl AsRef<str> for TransportKey {
59    fn as_ref(&self) -> &str {
60        &self.0
61    }
62}
63
64// From conversions for ergonomic construction
65impl From<&str> for TransportKey {
66    fn from(s: &str) -> Self {
67        Self(Arc::from(s))
68    }
69}
70
71impl From<String> for TransportKey {
72    fn from(s: String) -> Self {
73        Self(Arc::from(s))
74    }
75}
76
77impl From<Arc<str>> for TransportKey {
78    fn from(s: Arc<str>) -> Self {
79        Self(s)
80    }
81}
82
83impl From<&String> for TransportKey {
84    fn from(s: &String) -> Self {
85        Self(Arc::from(s.as_str()))
86    }
87}
88
89impl From<TransportKey> for String {
90    fn from(val: TransportKey) -> Self {
91        val.0.to_string()
92    }
93}
94
95// Borrow trait for HashMap lookups with &str
96impl std::borrow::Borrow<str> for TransportKey {
97    fn borrow(&self) -> &str {
98        &self.0
99    }
100}
101
102// Display for printing
103impl fmt::Display for TransportKey {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        write!(f, "{}", self.0)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use std::collections::{HashMap, HashSet};
113
114    #[test]
115    fn test_transport_key_creation() {
116        // Test new() method
117        let key1 = TransportKey::new("tcp");
118        assert_eq!(key1.as_str(), "tcp");
119
120        // Test From<&str>
121        let key2: TransportKey = "rdma".into();
122        assert_eq!(key2.as_str(), "rdma");
123
124        // Test From<String>
125        let key3 = TransportKey::from(String::from("udp"));
126        assert_eq!(key3.as_str(), "udp");
127
128        // Test From<&String>
129        let s = String::from("grpc");
130        let key4 = TransportKey::from(&s);
131        assert_eq!(key4.as_str(), "grpc");
132
133        // Test From<Arc<str>>
134        let arc_str: Arc<str> = Arc::from("http");
135        let key5 = TransportKey::from(arc_str);
136        assert_eq!(key5.as_str(), "http");
137    }
138
139    #[test]
140    fn test_transport_key_deref() {
141        let key = TransportKey::from("tcp");
142
143        // Deref to str methods should work
144        assert_eq!(key.len(), 3);
145        assert_eq!(key.chars().count(), 3);
146        assert!(key.starts_with("tc"));
147        assert!(key.ends_with("cp"));
148
149        // Can use str slicing through Deref
150        assert_eq!(&key[0..2], "tc");
151    }
152
153    #[test]
154    fn test_transport_key_as_ref() {
155        let key = TransportKey::from("tcp");
156
157        // AsRef<str> allows passing to functions expecting &str
158        fn takes_str_ref(s: &str) -> usize {
159            s.len()
160        }
161
162        assert_eq!(takes_str_ref(&key), 3);
163        assert_eq!(takes_str_ref(key.as_ref()), 3);
164    }
165
166    #[test]
167    fn test_transport_key_display() {
168        let key = TransportKey::from("tcp");
169        assert_eq!(format!("{}", key), "tcp");
170        assert_eq!(key.to_string(), "tcp");
171    }
172
173    #[test]
174    fn test_transport_key_debug() {
175        let key = TransportKey::from("tcp");
176        let debug_str = format!("{:?}", key);
177        assert!(debug_str.contains("TransportKey"));
178        assert!(debug_str.contains("tcp"));
179    }
180
181    #[test]
182    fn test_transport_key_equality() {
183        let key1 = TransportKey::from("tcp");
184        let key2 = TransportKey::from("tcp");
185        let key3 = TransportKey::from("rdma");
186
187        assert_eq!(key1, key2);
188        assert_ne!(key1, key3);
189
190        // Test with different source types
191        let key4: TransportKey = String::from("tcp").into();
192        assert_eq!(key1, key4);
193    }
194
195    #[test]
196    fn test_transport_key_ordering() {
197        let mut keys = [
198            TransportKey::from("udp"),
199            TransportKey::from("tcp"),
200            TransportKey::from("rdma"),
201            TransportKey::from("grpc"),
202        ];
203
204        keys.sort();
205
206        assert_eq!(keys[0], TransportKey::from("grpc"));
207        assert_eq!(keys[1], TransportKey::from("rdma"));
208        assert_eq!(keys[2], TransportKey::from("tcp"));
209        assert_eq!(keys[3], TransportKey::from("udp"));
210    }
211
212    #[test]
213    fn test_transport_key_hash() {
214        let mut set = HashSet::new();
215        set.insert(TransportKey::from("tcp"));
216        set.insert(TransportKey::from("rdma"));
217        set.insert(TransportKey::from("tcp")); // Duplicate
218
219        assert_eq!(set.len(), 2);
220        assert!(set.contains(&TransportKey::from("tcp")));
221        assert!(set.contains(&TransportKey::from("rdma")));
222        assert!(!set.contains(&TransportKey::from("udp")));
223    }
224
225    #[test]
226    fn test_transport_key_in_hashmap() {
227        let mut map = HashMap::new();
228        map.insert(TransportKey::from("tcp"), "tcp://127.0.0.1:5555");
229        map.insert(TransportKey::from("rdma"), "rdma://10.0.0.1:6666");
230
231        // Can lookup with TransportKey
232        assert_eq!(
233            map.get(&TransportKey::from("tcp")),
234            Some(&"tcp://127.0.0.1:5555")
235        );
236
237        // Can lookup with &str via Borrow trait
238        assert_eq!(map.get("tcp"), Some(&"tcp://127.0.0.1:5555"));
239        assert_eq!(map.get("rdma"), Some(&"rdma://10.0.0.1:6666"));
240        assert_eq!(map.get("udp"), None);
241    }
242
243    #[test]
244    fn test_transport_key_clone() {
245        let key1 = TransportKey::from("tcp");
246        let key2 = key1.clone();
247
248        assert_eq!(key1, key2);
249        assert_eq!(key1.as_str(), key2.as_str());
250
251        // Verify Arc is shared (same pointer)
252        let ptr1 = key1.as_str().as_ptr();
253        let ptr2 = key2.as_str().as_ptr();
254        assert_eq!(ptr1, ptr2);
255    }
256}
velo_ext/id/transport.rs

velo_ext/id/
transport.rs
