Skip to content

Add support for gRPC streams. #101

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
PiotrSikora merged 7 commits into from
May 19, 2021
Merged

Add support for gRPC streams. #101

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
1d9377b
gRPC stream support
Shikugawa Apr 13, 2021
0035f21
fix
Shikugawa May 5, 2021
f3fe656
fix with reviews
Shikugawa May 9, 2021
e1551b2
clippy
Shikugawa May 9, 2021
ac7f114
fix
Shikugawa May 15, 2021
8e9169b
fix
Shikugawa May 15, 2021
d2aaf0a
fix
Shikugawa May 15, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 114 additions & 14 deletions src/dispatcher.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,10 @@ pub(crate) fn register_grpc_callout(token_id: u32) {
DISPATCHER.with(|dispatcher| dispatcher.register_grpc_callout(token_id));
}

pub(crate) fn register_grpc_stream(token_id: u32) {
DISPATCHER.with(|dispatcher| dispatcher.register_grpc_stream(token_id));
}

struct NoopRoot;

impl Context for NoopRoot {}
Expand All @@ -57,6 +61,7 @@ struct Dispatcher {
active_id: Cell<u32>,
callouts: RefCell<HashMap<u32, u32>>,
grpc_callouts: RefCell<HashMap<u32, u32>>,
grpc_streams: RefCell<HashMap<u32, u32>>,
}

impl Dispatcher {
Expand All @@ -71,6 +76,7 @@ impl Dispatcher {
active_id: Cell::new(0),
callouts: RefCell::new(HashMap::new()),
grpc_callouts: RefCell::new(HashMap::new()),
grpc_streams: RefCell::new(HashMap::new()),
}
}

Expand All @@ -97,6 +103,17 @@ impl Dispatcher {
}
}

fn register_grpc_stream(&self, token_id: u32) {
if self
.grpc_streams
.borrow_mut()
.insert(token_id, self.active_id.get())
.is_some()
{
panic!("duplicate token_id")
}
}

fn register_grpc_callout(&self, token_id: u32) {
if self
.grpc_callouts
Expand Down Expand Up @@ -399,47 +416,116 @@ impl Dispatcher {
}
}

fn on_grpc_receive(&self, token_id: u32, response_size: usize) {
let context_id = self
.grpc_callouts
fn on_grpc_receive_initial_metadata(&self, token_id: u32, headers: u32) {
let context_id = *self
.grpc_streams
.borrow_mut()
.remove(&token_id)
.get(&token_id)
.expect("invalid token_id");

if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
http_stream.on_grpc_call_response(token_id, 0, response_size);
http_stream.on_grpc_stream_initial_metadata(token_id, headers);
} else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
stream.on_grpc_call_response(token_id, 0, response_size);
stream.on_grpc_stream_initial_metadata(token_id, headers);
} else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
root.on_grpc_call_response(token_id, 0, response_size);
root.on_grpc_stream_initial_metadata(token_id, headers);
}
}

fn on_grpc_close(&self, token_id: u32, status_code: u32) {
let context_id = self
.grpc_callouts
fn on_grpc_receive(&self, token_id: u32, response_size: usize) {
if let Some(context_id) = self.grpc_callouts.borrow_mut().remove(&token_id) {
if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
http_stream.on_grpc_call_response(token_id, 0, response_size);
} else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
stream.on_grpc_call_response(token_id, 0, response_size);
} else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
root.on_grpc_call_response(token_id, 0, response_size);
}
} else if let Some(context_id) = self.grpc_streams.borrow_mut().get(&token_id) {
let context_id = *context_id;
if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
http_stream.on_grpc_stream_message(token_id, response_size);
} else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
stream.on_grpc_stream_message(token_id, response_size);
} else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
root.on_grpc_stream_message(token_id, response_size);
}
} else {
panic!("invalid token_id")
}
}

fn on_grpc_receive_trailing_metadata(&self, token_id: u32, trailers: u32) {
let context_id = *self
.grpc_streams
.borrow_mut()
.remove(&token_id)
.get(&token_id)
.expect("invalid token_id");

if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
http_stream.on_grpc_call_response(token_id, status_code, 0);
http_stream.on_grpc_stream_initial_metadata(token_id, trailers);
Comment thread
Shikugawa marked this conversation as resolved.
Outdated
} else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
stream.on_grpc_call_response(token_id, status_code, 0);
stream.on_grpc_stream_initial_metadata(token_id, trailers);
Comment thread
Shikugawa marked this conversation as resolved.
Outdated
} else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
root.on_grpc_call_response(token_id, status_code, 0);
root.on_grpc_stream_initial_metadata(token_id, trailers);
Comment thread
Shikugawa marked this conversation as resolved.
Outdated
}
}

fn on_grpc_close(&self, token_id: u32, status_code: u32) {
if let Some(context_id) = self.grpc_callouts.borrow_mut().remove(&token_id) {
if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
http_stream.on_grpc_call_response(token_id, status_code, 0);
} else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
stream.on_grpc_call_response(token_id, status_code, 0);
} else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
root.on_grpc_call_response(token_id, status_code, 0);
}
} else if let Some(context_id) = self.grpc_streams.borrow_mut().remove(&token_id) {
if let Some(http_stream) = self.http_streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
http_stream.on_grpc_stream_close(token_id, status_code)
} else if let Some(stream) = self.streams.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
stream.on_grpc_stream_close(token_id, status_code)
} else if let Some(root) = self.roots.borrow_mut().get_mut(&context_id) {
self.active_id.set(context_id);
hostcalls::set_effective_context(context_id).unwrap();
root.on_grpc_stream_close(token_id, status_code)
}
} else {
panic!("invalid token_id")
}
}
}
Comment thread
PiotrSikora marked this conversation as resolved.
Expand Down Expand Up @@ -571,11 +657,25 @@ pub extern "C" fn proxy_on_http_call_response(
})
}

#[no_mangle]
pub extern "C" fn proxy_on_grpc_receive_initial_metadata(
_context_id: u32,
token_id: u32,
headers: u32,
) {
DISPATCHER.with(|dispatcher| dispatcher.on_grpc_receive_initial_metadata(token_id, headers))
}

#[no_mangle]
pub extern "C" fn proxy_on_grpc_receive(_context_id: u32, token_id: u32, response_size: usize) {
DISPATCHER.with(|dispatcher| dispatcher.on_grpc_receive(token_id, response_size))
}

#[no_mangle]
pub extern "C" fn proxy_on_grpc_trailing_metadata(_context_id: u32, token_id: u32, trailers: u32) {
Comment thread
Shikugawa marked this conversation as resolved.
Outdated
DISPATCHER.with(|dispatcher| dispatcher.on_grpc_receive_trailing_metadata(token_id, trailers))
}

#[no_mangle]
pub extern "C" fn proxy_on_grpc_close(_context_id: u32, token_id: u32, status_code: u32) {
DISPATCHER.with(|dispatcher| dispatcher.on_grpc_close(token_id, status_code))
Expand Down
131 changes: 129 additions & 2 deletions src/hostcalls.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,24 @@ pub fn get_map(map_type: MapType) -> Result<Vec<(String, String)>, Status> {
}
}

pub fn get_map_bytes(map_type: MapType) -> Result<Vec<(String, Vec<u8>)>, Status> {
unsafe {
let mut return_data: *mut u8 = null_mut();
let mut return_size: usize = 0;
match proxy_get_header_map_pairs(map_type, &mut return_data, &mut return_size) {
Status::Ok => {
if !return_data.is_null() {
let serialized_map = Vec::from_raw_parts(return_data, return_size, return_size);
Ok(utils::deserialize_bytes_map(&serialized_map))
} else {
Ok(Vec::new())
}
}
status => panic!("unexpected status: {}", status as u32),
}
}
}

extern "C" {
fn proxy_set_header_map_pairs(
map_type: MapType,
Expand Down Expand Up @@ -677,7 +695,7 @@ pub fn dispatch_grpc_call(
timeout: Duration,
) -> Result<u32, Status> {
let mut return_callout_id = 0;
let serialized_initial_metadata = utils::serialize_bytes_value_map(initial_metadata);
let serialized_initial_metadata = utils::serialize_bytes_map(initial_metadata);
unsafe {
match proxy_grpc_call(
upstream_name.as_ptr(),
Expand All @@ -704,6 +722,80 @@ pub fn dispatch_grpc_call(
}
}

extern "C" {
fn proxy_grpc_stream(
upstream_data: *const u8,
upstream_size: usize,
service_name_data: *const u8,
service_name_size: usize,
method_name_data: *const u8,
method_name_size: usize,
initial_metadata_data: *const u8,
initial_metadata_size: usize,
return_stream_id: *mut u32,
) -> Status;
}

pub fn open_grpc_stream(
upstream_name: &str,
service_name: &str,
method_name: &str,
initial_metadata: Vec<(&str, &[u8])>,
) -> Result<u32, Status> {
let mut return_stream_id = 0;
let serialized_initial_metadata = utils::serialize_bytes_map(initial_metadata);
unsafe {
match proxy_grpc_stream(
upstream_name.as_ptr(),
upstream_name.len(),
service_name.as_ptr(),
service_name.len(),
method_name.as_ptr(),
method_name.len(),
serialized_initial_metadata.as_ptr(),
serialized_initial_metadata.len(),
&mut return_stream_id,
) {
Status::Ok => {
dispatcher::register_grpc_stream(return_stream_id);
Ok(return_stream_id)
}
Status::ParseFailure => Err(Status::ParseFailure),
Status::InternalFailure => Err(Status::InternalFailure),
status => panic!("unexpected status: {}", status as u32),
}
}
}

extern "C" {
fn proxy_grpc_send(
token: u32,
message_ptr: *const u8,
message_len: usize,
end_stream: bool,
) -> Status;
}

pub fn send_grpc_stream_message(
token: u32,
message: Option<&[u8]>,
end_stream: bool,
) -> Result<(), Status> {
unsafe {
match proxy_grpc_send(
token,
message.map_or(null(), |message| message.as_ptr()),
message.map_or(0, |message| message.len()),
end_stream,
) {
Status::Ok => Ok(()),
Status::BadArgument => Err(Status::BadArgument),
Status::NotFound => Err(Status::NotFound),
status => panic!("unexpected status: {}", status as u32),
}
}
}

extern "C" {
fn proxy_grpc_cancel(token_id: u32) -> Status;
}
Expand All @@ -718,6 +810,20 @@ pub fn cancel_grpc_call(token_id: u32) -> Result<(), Status> {
}
}

extern "C" {
fn proxy_grpc_close(token_id: u32) -> Status;
}

pub fn close_grpc_stream(token_id: u32) -> Result<(), Status> {
unsafe {
match proxy_grpc_close(token_id) {
Status::Ok => Ok(()),
Status::NotFound => Err(Status::NotFound),
status => panic!("unexpected status: {}", status as u32),
}
}
}

extern "C" {
fn proxy_set_effective_context(context_id: u32) -> Status;
}
Expand Down Expand Up @@ -850,7 +956,7 @@ mod utils {
bytes
}

pub(super) fn serialize_bytes_value_map(map: Vec<(&str, &[u8])>) -> Bytes {
pub(super) fn serialize_bytes_map(map: Vec<(&str, &[u8])>) -> Bytes {
let mut size: usize = 4;
for (name, value) in &map {
size += name.len() + value.len() + 10;
Expand Down Expand Up @@ -893,4 +999,25 @@ mod utils {
}
map
}

pub(super) fn deserialize_bytes_map(bytes: &[u8]) -> Vec<(String, Vec<u8>)> {
Comment thread
PiotrSikora marked this conversation as resolved.
let mut map = Vec::new();
if bytes.is_empty() {
return map;
}
let size = u32::from_le_bytes(<[u8; 4]>::try_from(&bytes[0..4]).unwrap()) as usize;
let mut p = 4 + size * 8;
for n in 0..size {
let s = 4 + n * 8;
let size = u32::from_le_bytes(<[u8; 4]>::try_from(&bytes[s..s + 4]).unwrap()) as usize;
let key = bytes[p..p + size].to_vec();
p += size + 1;
let size =
u32::from_le_bytes(<[u8; 4]>::try_from(&bytes[s + 4..s + 8]).unwrap()) as usize;
let value = bytes[p..p + size].to_vec();
p += size + 1;
map.push((String::from_utf8(key).unwrap(), value));
}
map
}
}
Loading