ttrpc/
context.rs

1// Copyright (c) 2021 Ant group
2//
3// SPDX-License-Identifier: Apache-2.0
4//
5
6use crate::proto::KeyValue;
7use core::time::Duration;
8use std::collections::HashMap;
9#[derive(Clone, Default, Debug)]
10pub struct Context {
11    pub metadata: HashMap<String, Vec<String>>,
12    pub timeout_nano: i64,
13}
14
15pub fn with_timeout(i: i64) -> Context {
16    Context {
17        timeout_nano: i,
18        ..Default::default()
19    }
20}
21///With_timeout is not friendly enough to external interfaces
22///Keep with_timeout for forward compatibility
23pub fn with_duration(du: Duration) -> Context {
24    with_timeout(du.as_nanos() as i64)
25}
26
27pub fn with_metadata(md: HashMap<String, Vec<String>>) -> Context {
28    Context {
29        metadata: md,
30        ..Default::default()
31    }
32}
33
34impl Context {
35    // appends additional values to the given key.
36    pub fn add(&mut self, key: String, value: String) {
37        if let Some(ref mut vl) = self.metadata.get_mut(&key) {
38            vl.push(value);
39        } else {
40            self.metadata.insert(key.to_lowercase(), vec![value]);
41        }
42    }
43
44    // Set sets the provided values for a given key.
45    // The values will overwrite any existing values.
46    // If no values provided, a key will be deleted.
47    pub fn set(&mut self, key: String, value: Vec<String>) {
48        if value.is_empty() {
49            self.metadata.remove(&key);
50        } else {
51            self.metadata.insert(key.to_lowercase(), value);
52        }
53    }
54}
55
56pub fn from_pb(kvs: &Vec<KeyValue>) -> HashMap<String, Vec<String>> {
57    let mut meta: HashMap<String, Vec<String>> = HashMap::new();
58    for kv in kvs {
59        if let Some(ref mut vl) = meta.get_mut(&kv.key) {
60            vl.push(kv.value.clone());
61        } else {
62            meta.insert(kv.key.clone(), vec![kv.value.clone()]);
63        }
64    }
65    meta
66}
67
68pub fn to_pb(kvs: HashMap<String, Vec<String>>) -> Vec<KeyValue> {
69    let mut meta = Vec::with_capacity(kvs.len());
70
71    for (k, vl) in kvs {
72        for v in vl {
73            let key = KeyValue {
74                key: k.clone(),
75                value: v.clone(),
76                ..Default::default()
77            };
78            meta.push(key);
79        }
80    }
81
82    meta
83}
84
85#[cfg(test)]
86mod tests {
87    use crate::context;
88    use crate::proto::KeyValue;
89
90    #[test]
91    fn test_metadata() {
92        // RepeatedField -> HashMap, test from_pb()
93        let mut src = Vec::new();
94        for i in &[
95            ("key1", "value1-1"),
96            ("key1", "value1-2"),
97            ("key2", "value2"),
98        ] {
99            let key = KeyValue {
100                key: i.0.to_string(),
101                value: i.1.to_string(),
102                ..Default::default()
103            };
104            src.push(key);
105        }
106
107        let dst = context::from_pb(&src);
108        assert_eq!(dst.len(), 2);
109
110        assert_eq!(
111            dst.get("key1"),
112            Some(&vec!["value1-1".to_string(), "value1-2".to_string()])
113        );
114        assert_eq!(dst.get("key2"), Some(&vec!["value2".to_string()]));
115        assert_eq!(dst.get("key3"), None);
116
117        // HashMap -> RepeatedField , test to_pb()
118        let mut kvs = context::to_pb(dst);
119        kvs.sort_by(|a, b| a.key.partial_cmp(&b.key).unwrap());
120
121        assert_eq!(kvs.len(), 3);
122
123        assert_eq!(kvs[0].key, "key1");
124        assert_eq!(kvs[0].value, "value1-1");
125
126        assert_eq!(kvs[1].key, "key1");
127        assert_eq!(kvs[1].value, "value1-2");
128
129        assert_eq!(kvs[2].key, "key2");
130        assert_eq!(kvs[2].value, "value2");
131    }
132
133    #[test]
134    fn test_context() {
135        let ctx: context::Context = Default::default();
136        assert_eq!(0, ctx.timeout_nano);
137        assert_eq!(ctx.metadata.len(), 0);
138
139        let mut ctx = context::with_duration(core::time::Duration::from_nanos(99));
140        assert_eq!(99, ctx.timeout_nano);
141        assert_eq!(ctx.metadata.len(), 0);
142
143        ctx.add("key1".to_string(), "value1-1".to_string());
144        assert_eq!(ctx.metadata.len(), 1);
145        assert_eq!(
146            ctx.metadata.get("key1"),
147            Some(&vec!["value1-1".to_string()])
148        );
149
150        ctx.add("key1".to_string(), "value1-2".to_string());
151        assert_eq!(ctx.metadata.len(), 1);
152        assert_eq!(
153            ctx.metadata.get("key1"),
154            Some(&vec!["value1-1".to_string(), "value1-2".to_string()])
155        );
156
157        ctx.set("key2".to_string(), vec!["value2".to_string()]);
158        assert_eq!(ctx.metadata.len(), 2);
159        assert_eq!(ctx.metadata.get("key2"), Some(&vec!["value2".to_string()]));
160
161        ctx.set("key1".to_string(), vec![]);
162        assert_eq!(ctx.metadata.len(), 1);
163        assert_eq!(ctx.metadata.get("key1"), None);
164    }
165}
ttrpc/context.rs

ttrpc/
context.rs
