Skip to content

Commit

Permalink
robustness: remove unwrap to prevent potential panics
Browse files Browse the repository at this point in the history
  • Loading branch information
wanglei13866 committed Sep 30, 2024
1 parent 3f4274e commit 57e8281
Show file tree
Hide file tree
Showing 3 changed files with 740 additions and 203 deletions.
107 changes: 85 additions & 22 deletions src/dispatcher.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use crate::hostcalls;
use crate::traits::*;
use crate::types::*;
use hashbrown::HashMap;
use log::trace;
use log::{error, trace};
use std::cell::{Cell, RefCell};

thread_local! {
Expand Down Expand Up @@ -414,15 +414,24 @@ impl Dispatcher {

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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
http_stream.on_http_call_response(token_id, num_headers, body_size, num_trailers)
} 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();
if let Err(e) = hostcalls::set_effective_context(context_id) {
error!("set_effective_context failed: {:?}", e);
return;
}
stream.on_http_call_response(token_id, num_headers, body_size, num_trailers)
} 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();
if let Err(e) = hostcalls::set_effective_context(context_id) {
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_http_call_response(token_id, num_headers, body_size, num_trailers)
}
}
Expand All @@ -439,15 +448,24 @@ impl Dispatcher {

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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_grpc_stream_initial_metadata(token_id, headers);
}
}
Expand All @@ -457,31 +475,49 @@ impl Dispatcher {
if let Some(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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id) {
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_grpc_call_response(token_id, 0, response_size);
}
} else {
let context_id = self.grpc_streams.borrow().get(&token_id).cloned();
if let Some(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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_grpc_stream_message(token_id, response_size);
}
} else {
Expand All @@ -503,15 +539,24 @@ impl Dispatcher {

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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
http_stream.on_grpc_stream_trailing_metadata(token_id, trailers);
} 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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
stream.on_grpc_stream_trailing_metadata(token_id, trailers);
} 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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_grpc_stream_trailing_metadata(token_id, trailers);
}
}
Expand All @@ -521,31 +566,49 @@ impl Dispatcher {
if let Some(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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_grpc_call_response(token_id, status_code, 0);
}
} else {
let context_id = self.grpc_streams.borrow_mut().remove(&token_id);
if let Some(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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
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();
if let Err(e) = hostcalls::set_effective_context(context_id){
error!("set_effective_context failed: {:?}", e);
return;
}
root.on_grpc_stream_close(token_id, status_code)
}
} else {
Expand Down
Loading

0 comments on commit 57e8281

Please sign in to comment.