Skip to content

Commit f7d7708

Browse files
committed
api: add metadata for ttrpc context
For tracing purpose, it needs to pass metadata like span id to the server. `ttrpc-rust` should then parse the metadata to service handler by `TtrpcContext` This commit is mainly a porting of parts of `metadata_test.go` and `types.go` if go version, and now only working on server side. Signed-off-by: bin <bin@hyper.sh>
1 parent 72db426 commit f7d7708

File tree

6 files changed

+57
-1
lines changed

6 files changed

+57
-1
lines changed

src/asynchronous/server.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -310,7 +310,7 @@ async fn handle_request(
310310
let path = format!("/{}/{}", req.service, req.method);
311311
if let Some(x) = methods.get(&path) {
312312
let method = x;
313-
let ctx = TtrpcContext { fd, mh: header };
313+
let ctx = TtrpcContext { fd, mh: header, metadata: common::parse_metadata(&req.metadata) };
314314

315315
match method.handler(ctx, req).await {
316316
Ok((stream_id, body)) => {

src/asynchronous/utils.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use crate::error::{Error, Result};
88
use crate::ttrpc::{Request, Response};
99
use async_trait::async_trait;
1010
use protobuf::Message;
11+
use std::collections::HashMap;
1112
use std::os::unix::io::{FromRawFd, RawFd};
1213
use tokio::net::UnixStream;
1314

@@ -87,6 +88,7 @@ pub trait MethodHandler {
8788
pub struct TtrpcContext {
8889
pub fd: std::os::unix::io::RawFd,
8990
pub mh: MessageHeader,
91+
pub metadata: HashMap<String, Vec<String>>,
9092
}
9193

9294
pub fn convert_response_to_buf(res: Response) -> Result<Vec<u8>> {

src/common.rs

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,10 @@
88
#![allow(unused_macros)]
99

1010
use crate::error::{Error, Result};
11+
use crate::ttrpc::KeyValue;
1112
use nix::fcntl::{fcntl, FcntlArg, FdFlag, OFlag};
1213
use nix::sys::socket::*;
14+
use std::collections::HashMap;
1315
use std::os::unix::io::RawFd;
1416
use std::str::FromStr;
1517

@@ -139,3 +141,46 @@ pub const MESSAGE_LENGTH_MAX: usize = 4 << 20;
139141

140142
pub const MESSAGE_TYPE_REQUEST: u8 = 0x1;
141143
pub const MESSAGE_TYPE_RESPONSE: u8 = 0x2;
144+
145+
pub fn parse_metadata(kvs: &protobuf::RepeatedField<KeyValue>) -> HashMap<String, Vec<String>> {
146+
let mut meta: HashMap<String, Vec<String>> = HashMap::new();
147+
for kv in kvs {
148+
if let Some(ref mut vl) = meta.get_mut(&kv.key) {
149+
vl.push(kv.value.clone());
150+
} else {
151+
meta.insert(kv.key.clone(), vec![kv.value.clone()]);
152+
}
153+
}
154+
meta
155+
}
156+
157+
#[cfg(test)]
158+
mod tests {
159+
use super::parse_metadata;
160+
use crate::ttrpc::KeyValue;
161+
162+
#[test]
163+
fn test_parse_metadata() {
164+
let mut src: protobuf::RepeatedField<KeyValue> = protobuf::RepeatedField::default();
165+
for i in &[
166+
("key1", "value1-1"),
167+
("key1", "value1-2"),
168+
("key2", "value2"),
169+
] {
170+
let mut key = KeyValue::default();
171+
key.key = i.0.to_string();
172+
key.value = i.1.to_string();
173+
src.push(key);
174+
}
175+
176+
let dst = parse_metadata(&src);
177+
assert_eq!(dst.len(), 2);
178+
179+
assert_eq!(
180+
dst.get("key1"),
181+
Some(&vec!["value1-1".to_string(), "value1-2".to_string()])
182+
);
183+
assert_eq!(dst.get("key2"), Some(&vec!["value2".to_string()]));
184+
assert_eq!(dst.get("key3"), None);
185+
}
186+
}

src/sync/server.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ fn start_method_handler_thread(
209209
fd,
210210
mh,
211211
res_tx: res_tx.clone(),
212+
metadata: common::parse_metadata(&req.metadata),
212213
};
213214
if let Err(x) = method.handler(ctx, req) {
214215
debug!("method handle {} get error {:?}", path, x);

src/sync/utils.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ use crate::common::{MessageHeader, MESSAGE_TYPE_RESPONSE};
77
use crate::error::{Error, Result};
88
use crate::ttrpc::{Request, Response};
99
use protobuf::Message;
10+
use std::collections::HashMap;
1011

1112
/// Response message through a channel.
1213
/// Eventually the message will sent to Client.
@@ -94,6 +95,7 @@ pub struct TtrpcContext {
9495
pub fd: std::os::unix::io::RawFd,
9596
pub mh: MessageHeader,
9697
pub res_tx: std::sync::mpsc::Sender<(MessageHeader, Vec<u8>)>,
98+
pub metadata: HashMap<String, Vec<String>>,
9799
}
98100

99101
/// Trait that implements handler which is a proxy to the desired method (sync).

src/ttrpc.proto

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,12 @@ message Request {
2121
string method = 2;
2222
bytes payload = 3;
2323
int64 timeout_nano = 4;
24+
repeated KeyValue metadata = 5;
25+
}
26+
27+
message KeyValue {
28+
string key = 1;
29+
string value = 2;
2430
}
2531

2632
// Get from github.com/gogo/protobuf/protobuf/google/protobuf/any.proto

0 commit comments

Comments
 (0)