protocols.rs 9.76 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
Ryan Olson's avatar
Ryan Olson committed
15
16

use serde::{Deserialize, Serialize};
17
18
19
use std::str::FromStr;

use crate::pipeline::PipelineError;
Ryan Olson's avatar
Ryan Olson committed
20
21
22

pub mod annotated;

23
24
pub type LeaseId = i64;

25
26
27
28
29
30
31
/// Default namespace if user does not provide one
const DEFAULT_NAMESPACE: &str = "NS";

const DEFAULT_COMPONENT: &str = "C";

const DEFAULT_ENDPOINT: &str = "E";

32
33
34
35
36
/// How we identify a namespace/component/endpoint URL.
/// Technically the '://' is not part of the scheme but it eliminates several string
/// concatenations.
pub const ENDPOINT_SCHEME: &str = "dyn://";

Ryan Olson's avatar
Ryan Olson committed
37
38
39
40
41
42
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct Component {
    pub name: String,
    pub namespace: String,
}

43
44
/// Represents an endpoint with a namespace, component, and name.
///
45
/// An `Endpoint` is defined by a three-part string separated by `/` or a '.':
46
47
48
49
50
/// - **namespace**
/// - **component**
/// - **name**
///
/// Example format: `"namespace/component/endpoint"`
51
52
///
/// TODO: There is also an Endpoint in runtime/src/component.rs
Ryan Olson's avatar
Ryan Olson committed
53
54
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct Endpoint {
55
56
    pub namespace: String,
    pub component: String,
Ryan Olson's avatar
Ryan Olson committed
57
    pub name: String,
58
}
59

60
61
62
63
64
impl PartialEq<Vec<&str>> for Endpoint {
    fn eq(&self, other: &Vec<&str>) -> bool {
        if other.len() != 3 {
            return false;
        }
65

66
67
68
69
70
71
72
73
        self.namespace == other[0] && self.component == other[1] && self.name == other[2]
    }
}

impl PartialEq<Endpoint> for Vec<&str> {
    fn eq(&self, other: &Endpoint) -> bool {
        other == self
    }
Ryan Olson's avatar
Ryan Olson committed
74
75
}

76
77
78
79
80
81
82
83
84
impl Default for Endpoint {
    fn default() -> Self {
        Endpoint {
            namespace: DEFAULT_NAMESPACE.to_string(),
            component: DEFAULT_COMPONENT.to_string(),
            name: DEFAULT_ENDPOINT.to_string(),
        }
    }
}
85

86
87
impl From<&str> for Endpoint {
    /// Creates an `Endpoint` from a string.
88
89
90
91
    ///
    /// # Arguments
    /// - `path`: A string in the format `"namespace/component/endpoint"`.
    ///
92
93
94
95
96
97
98
99
100
101
    /// The first two parts become the first two elements of the vector.
    /// The third and subsequent parts are joined with '_' and become the third element.
    /// Default values are used for missing parts.
    ///
    /// # Examples:
    /// - "component" -> ["DEFAULT_NS", "component", "DEFAULT_E"]
    /// - "namespace.component" -> ["namespace", "component", "DEFAULT_E"]
    /// - "namespace.component.endpoint" -> ["namespace", "component", "endpoint"]
    /// - "namespace/component" -> ["namespace", "component", "DEFAULT_E"]
    /// - "namespace.component.endpoint.other.parts" -> ["namespace", "component", "endpoint_other_parts"]
102
103
104
    ///
    /// # Examples
    /// ```ignore
Neelay Shah's avatar
Neelay Shah committed
105
    /// use dynamo_runtime:protocols::Endpoint;
106
    ///
107
    /// let endpoint = Endpoint::from("namespace/component/endpoint");
108
109
110
111
    /// assert_eq!(endpoint.namespace, "namespace");
    /// assert_eq!(endpoint.component, "component");
    /// assert_eq!(endpoint.name, "endpoint");
    /// ```
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
    fn from(input: &str) -> Self {
        let mut result = Endpoint::default();

        // Split the input string on either '.' or '/'
        let elements: Vec<&str> = input
            .trim_matches([' ', '/', '.'])
            .split(['.', '/'])
            .filter(|x| !x.is_empty())
            .collect();

        match elements.len() {
            0 => {}
            1 => {
                result.component = elements[0].to_string();
            }
            2 => {
                result.namespace = elements[0].to_string();
                result.component = elements[1].to_string();
            }
            3 => {
                result.namespace = elements[0].to_string();
                result.component = elements[1].to_string();
                result.name = elements[2].to_string();
            }
            x if x > 3 => {
                result.namespace = elements[0].to_string();
                result.component = elements[1].to_string();
                result.name = elements[2..].join("_");
            }
            _ => unreachable!(),
142
        }
143
        result
144
145
146
147
148
149
150
151
    }
}

impl FromStr for Endpoint {
    type Err = PipelineError;

    /// Parses an `Endpoint` from a string using the standard Rust `.parse::<T>()` pattern.
    ///
152
    /// This is implemented in terms of [`From<&str>`].
153
154
    ///
    /// # Errors
155
    /// Does not fail
156
157
158
159
    ///
    /// # Examples
    /// ```ignore
    /// use std::str::FromStr;
Neelay Shah's avatar
Neelay Shah committed
160
    /// use dynamo_runtime:protocols::Endpoint;
161
162
163
164
165
    ///
    /// let endpoint: Endpoint = "namespace/component/endpoint".parse().unwrap();
    /// assert_eq!(endpoint.namespace, "namespace");
    /// assert_eq!(endpoint.component, "component");
    /// assert_eq!(endpoint.name, "endpoint");
166
167
168
    /// let endpoint: Endpoint = "dyn://namespace/component/endpoint".parse().unwrap();
    /// // same as above
    /// assert_eq!(endpoint.name, "endpoint");
169
170
    /// ```
    fn from_str(s: &str) -> Result<Self, Self::Err> {
171
172
        let cleaned = s.strip_prefix(ENDPOINT_SCHEME).unwrap_or(s);
        Ok(Endpoint::from(cleaned))
173
174
175
    }
}

176
177
178
179
180
181
182
183
184
185
impl Endpoint {
    /// As a String like dyn://dynamo.internal.worker
    pub fn as_url(&self) -> String {
        format!(
            "{ENDPOINT_SCHEME}{}.{}.{}",
            self.namespace, self.component, self.name
        )
    }
}

Ryan Olson's avatar
Ryan Olson committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum RouterType {
    PushRoundRobin,
    PushRandom,
}

impl Default for RouterType {
    fn default() -> Self {
        Self::PushRandom
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub struct ModelMetaData {
    pub name: String,
    pub component: Component,
    pub router_type: RouterType,
}
205
206
207
208

#[cfg(test)]
mod tests {
    use super::*;
209
210
    use std::convert::TryFrom;
    use std::str::FromStr;
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239

    #[test]
    fn test_router_type_default() {
        let default_router = RouterType::default();
        assert_eq!(default_router, RouterType::PushRandom);
    }

    #[test]
    fn test_router_type_serialization() {
        let router_round_robin = RouterType::PushRoundRobin;
        let router_random = RouterType::PushRandom;

        let serialized_round_robin = serde_json::to_string(&router_round_robin).unwrap();
        let serialized_random = serde_json::to_string(&router_random).unwrap();

        assert_eq!(serialized_round_robin, "\"push_round_robin\"");
        assert_eq!(serialized_random, "\"push_random\"");
    }

    #[test]
    fn test_router_type_deserialization() {
        let round_robin: RouterType = serde_json::from_str("\"push_round_robin\"").unwrap();
        let random: RouterType = serde_json::from_str("\"push_random\"").unwrap();

        assert_eq!(round_robin, RouterType::PushRoundRobin);
        assert_eq!(random, RouterType::PushRandom);
    }

    #[test]
240
    fn test_valid_endpoint_from() {
241
        let input = "namespace1/component1/endpoint1";
242
        let endpoint = Endpoint::from(input);
243
244
245
246
247
248
249
250
251

        assert_eq!(endpoint.namespace, "namespace1");
        assert_eq!(endpoint.component, "component1");
        assert_eq!(endpoint.name, "endpoint1");
    }

    #[test]
    fn test_valid_endpoint_from_str() {
        let input = "namespace2/component2/endpoint2";
252
        let endpoint = Endpoint::from_str(input).unwrap();
253
254
255
256
257
258
259
260
261

        assert_eq!(endpoint.namespace, "namespace2");
        assert_eq!(endpoint.component, "component2");
        assert_eq!(endpoint.name, "endpoint2");
    }

    #[test]
    fn test_valid_endpoint_parse() {
        let input = "namespace3/component3/endpoint3";
262
        let endpoint: Endpoint = input.parse().unwrap();
263
264
265
266
267
268
269

        assert_eq!(endpoint.namespace, "namespace3");
        assert_eq!(endpoint.component, "component3");
        assert_eq!(endpoint.name, "endpoint3");
    }

    #[test]
270
271
    fn test_endpoint_from() {
        let result = Endpoint::from("component");
272
        assert_eq!(
273
274
            result,
            vec![DEFAULT_NAMESPACE, "component", DEFAULT_ENDPOINT]
275
276
277
278
        );
    }

    #[test]
279
280
281
    fn test_namespace_component_endpoint() {
        let result = Endpoint::from("namespace.component.endpoint");
        assert_eq!(result, vec!["namespace", "component", "endpoint"]);
282
283
284
    }

    #[test]
285
286
287
    fn test_forward_slash_separator() {
        let result = Endpoint::from("namespace/component");
        assert_eq!(result, vec!["namespace", "component", DEFAULT_ENDPOINT]);
288
289
290
    }

    #[test]
291
292
    fn test_multiple_parts() {
        let result = Endpoint::from("namespace.component.endpoint.other.parts");
293
        assert_eq!(
294
295
            result,
            vec!["namespace", "component", "endpoint_other_parts"]
296
297
298
299
        );
    }

    #[test]
300
301
302
303
    fn test_mixed_separators() {
        // Do it the .into way for variety and documentation
        let result: Endpoint = "namespace/component.endpoint".into();
        assert_eq!(result, vec!["namespace", "component", "endpoint"]);
304
305
306
    }

    #[test]
307
308
309
310
311
312
    fn test_empty_string() {
        let result = Endpoint::from("");
        assert_eq!(
            result,
            vec![DEFAULT_NAMESPACE, DEFAULT_COMPONENT, DEFAULT_ENDPOINT]
        );
313

314
315
316
317
318
        // White space is equivalent to an empty string
        let result = Endpoint::from("   ");
        assert_eq!(
            result,
            vec![DEFAULT_NAMESPACE, DEFAULT_COMPONENT, DEFAULT_ENDPOINT]
319
320
        );
    }
321
}