fix with reviews

Signed-off-by: Shikugawa <[email protected]>

Author: Shikugawa

src/dispatcher.rs

@@ -61,7 +61,7 @@ struct Dispatcher {
     active_id: Cell<u32>,
     callouts: RefCell<HashMap<u32, u32>>,
     grpc_callouts: RefCell<HashMap<u32, u32>>,
-    grpc_stream_tokens: RefCell<HashMap<u32, u32>>,
+    grpc_streams: RefCell<HashMap<u32, u32>>,
 }
 
 impl Dispatcher {
@@ -76,7 +76,7 @@ impl Dispatcher {
             active_id: Cell::new(0),
             callouts: RefCell::new(HashMap::new()),
             grpc_callouts: RefCell::new(HashMap::new()),
-            grpc_stream_tokens: RefCell::new(HashMap::new()),
+            grpc_streams: RefCell::new(HashMap::new()),
         }
     }
 
@@ -103,9 +103,9 @@ impl Dispatcher {
         }
     }
 
-    fn register_grpc_stream_tokens(&self, token_id: u32) {
+    fn register_grpc_stream(&self, token_id: u32) {
         if self
-            .grpc_stream_tokens
+            .grpc_streams
             .borrow_mut()
             .insert(token_id, self.active_id.get())
             .is_some()
@@ -416,6 +416,29 @@ impl Dispatcher {
         }
     }
 
+    fn on_grpc_receive_initial_metadata(&self, token_id: u32, headers: u32) {
+        let context_id = self
+            .grpc_streams
+            .borrow_mut()
+            .get(&token_id)
+            .expect("invalid token_id")
+            .clone();
+
+        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_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_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_stream_initial_metadata(token_id, headers);
+        }
+    }
+
     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) {
@@ -431,26 +454,49 @@ impl Dispatcher {
                 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_callouts.borrow_mut().get(&token_id) {
+        } else if let Some(context_id) = self.grpc_streams.borrow_mut().get(&token_id) {
             let context_id = context_id.clone();
             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_receive_body(token_id, response_size);
+                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_receive_body(token_id, response_size);
+                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_receive_body(token_id, response_size);
+                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()
+            .get(&token_id)
+            .expect("invalid token_id")
+            .clone();
+
+        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_initial_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();
+            stream.on_grpc_stream_initial_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();
+            root.on_grpc_stream_initial_metadata(token_id, trailers);
+        }
+    }
+
     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) {
@@ -466,7 +512,7 @@ impl Dispatcher {
                 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_stream_tokens.borrow_mut().remove(&token_id) {
+        } 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();
@@ -484,52 +530,6 @@ impl Dispatcher {
             panic!("invalid token_id")
         }
     }
-
-    fn on_grpc_receive_initial_metadata(&self, token_id: u32, headers: u32) {
-        let context_id = self
-            .grpc_stream_tokens
-            .borrow_mut()
-            .get(&token_id)
-            .expect("invalid token_id")
-            .clone();
-
-        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_receive_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_stream_receive_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_stream_receive_initial_metadata(token_id, headers);
-        }
-    }
-
-    fn on_grpc_receive_trailing_metadata(&self, token_id: u32, trailers: u32) {
-        let context_id = self
-            .grpc_stream_tokens
-            .borrow_mut()
-            .get(&token_id)
-            .expect("invalid token_id")
-            .clone();
-
-        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_receive_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();
-            stream.on_grpc_stream_receive_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();
-            root.on_grpc_stream_receive_trailing_metadata(token_id, trailers);
-        }
-    }
 }
 
 #[no_mangle]
@@ -659,11 +659,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) {
+    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))

src/hostcalls.rs

@@ -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,
@@ -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(),
@@ -718,14 +736,14 @@ extern "C" {
     ) -> Status;
 }
 
-pub fn create_grpc_stream(
+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_value_map(initial_metadata);
+    let serialized_initial_metadata = utils::serialize_bytes_map(initial_metadata);
     unsafe {
         match proxy_grpc_stream(
             upstream_name.as_ptr(),
@@ -758,7 +776,7 @@ extern "C" {
     ) -> Status;
 }
 
-pub fn grpc_stream_send(
+pub fn send_grpc_stream_message(
     token: u32,
     message: Option<&[u8]>,
     end_stream: bool,
@@ -782,7 +800,7 @@ extern "C" {
     fn proxy_grpc_cancel(token_id: u32) -> Status;
 }
 
-pub fn grpc_call_cancel(token_id: u32) -> Result<(), Status> {
+pub fn cancel_grpc_call(token_id: u32) -> Result<(), Status> {
     unsafe {
         match proxy_grpc_cancel(token_id) {
             Status::Ok => Ok(()),
@@ -796,7 +814,7 @@ extern "C" {
     fn proxy_grpc_close(token_id: u32) -> Status;
 }
 
-pub fn grpc_stream_close(token_id: u32) -> Result<(), Status> {
+pub fn close_grpc_stream(token_id: u32) -> Result<(), Status> {
     unsafe {
         match proxy_grpc_close(token_id) {
             Status::Ok => Ok(()),
@@ -938,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;

src/traits.rs

@@ -119,45 +119,49 @@ pub trait Context {
         hostcalls::cancel_grpc_call(token_id)
     }
 
-    fn create_grpc_stream(
+    fn open_grpc_stream(
         &self,
         cluster_name: &str,
         service_name: &str,
         method_name: &str,
         initial_metadata: Vec<(&str, &[u8])>,
     ) -> Result<u32, Status> {
-        hostcalls::create_grpc_stream(cluster_name, service_name, method_name, initial_metadata)
+        hostcalls::open_grpc_stream(cluster_name, service_name, method_name, initial_metadata)
     }
 
-    fn grpc_stream_send(
+    fn on_grpc_stream_initial_metadata(&mut self, _token_id: u32, _num_elements: u32) {}
+
+    fn get_grpc_stream_initial_metadata(&self) -> Vec<(String, Vec<u8>)> {
+        hostcalls::get_map_bytes(MapType::GrpcReceiveInitialMetadata).unwrap()
+    }
+
+    fn send_grpc_stream_message(
         &self,
         token_id: u32,
         message: Option<&[u8]>,
         end_stream: bool,
     ) -> Result<(), Status> {
-        hostcalls::grpc_stream_send(token_id, message, end_stream)
+        hostcalls::send_grpc_stream_message(token_id, message, end_stream)
     }
 
-    fn grpc_stream_close(&self, token_id: u32) -> Result<(), Status> {
-        hostcalls::grpc_stream_close(token_id)
-    }
-
-    fn on_grpc_stream_receive_initial_metadata(&mut self, _token_id: u32, _headers: u32) {}
+    fn on_grpc_stream_message(&mut self, _token_id: u32, _message_size: usize) {}
 
-    fn on_grpc_stream_receive_trailing_metadata(&mut self, _token_id: u32, _trailers: u32) {}
-
-    fn on_grpc_stream_receive_body(&mut self, _token_id: u32, _response_size: usize) {}
+    fn get_grpc_stream_message(&mut self, start: usize, max_size: usize) -> Option<Bytes> {
+        hostcalls::get_buffer(BufferType::GrpcReceiveBuffer, start, max_size).unwrap()
+    }
 
-    fn on_grpc_stream_close(&mut self, _token_id: u32, _status_code: u32) {}
+    fn on_grpc_stream_trailing_metadata(&mut self, _token_id: u32, _num_elements: u32) {}
 
-    fn get_grpc_call_initial_metadata(&self) -> Vec<(String, Vec<u8>)> {
-        hostcalls::get_map_bytes(MapType::GrpcReceiveInitialMetadata).unwrap()
+    fn get_grpc_stream_trailing_metadata(&self) -> Vec<(String, Vec<u8>)> {
+        hostcalls::get_map_bytes(MapType::GrpcReceiveTrailingMetadata).unwrap()
     }
 
-    fn get_grpc_call_trailing_metadata(&self) -> Vec<(String, Vec<u8>)> {
-        hostcalls::get_map_bytes(MapType::GrpcReceiveTrailingMetadata).unwrap()
+    fn close_grpc_stream(&self, token_id: u32) -> Result<(), Status> {
+        hostcalls::close_grpc_stream(token_id)
     }
 
+    fn on_grpc_stream_close(&mut self, _token_id: u32, _status_code: u32) {}
+
     fn on_done(&mut self) -> bool {
         true
     }