diff --git a/src/asynchronous/server.rs b/src/asynchronous/server.rs
index 26c49f2c..9158e648 100644
--- a/src/asynchronous/server.rs
+++ b/src/asynchronous/server.rs
@@ -381,12 +381,14 @@ impl ReaderDelegate for ServerReader {
async fn handle_msg(&self, msg: GenMessage) {
let handler_shutdown_waiter = self.handler_shutdown.subscribe();
let context = self.context();
+ let (wait_tx, wait_rx) = tokio::sync::oneshot::channel::<()>();
spawn(async move {
select! {
- _ = context.handle_msg(msg) => {}
+ _ = context.handle_msg(msg, wait_tx) => {}
_ = handler_shutdown_waiter.wait_shutdown() => {}
}
});
+ wait_rx.await.unwrap_or_default();
}
async fn handle_err(&self, header: MessageHeader, e: Error) {
@@ -424,7 +426,7 @@ impl HandlerContext {
})
.ok();
}
- async fn handle_msg(&self, msg: GenMessage) {
+ async fn handle_msg(&self, msg: GenMessage, wait_tx: tokio::sync::oneshot::Sender<()>) {
let stream_id = msg.header.stream_id;
if (stream_id % 2) != 1 {
@@ -438,7 +440,7 @@ impl HandlerContext {
}
match msg.header.type_ {
- MESSAGE_TYPE_REQUEST => match self.handle_request(msg).await {
+ MESSAGE_TYPE_REQUEST => match self.handle_request(msg, wait_tx).await {
Ok(opt_msg) => match opt_msg {
Some(mut resp) => {
// Server: check size before sending to client
@@ -471,6 +473,8 @@ impl HandlerContext {
Err(status) => Self::respond_with_status(self.tx.clone(), stream_id, status).await,
},
MESSAGE_TYPE_DATA => {
+ // no need to wait data message handling
+ drop(wait_tx);
// TODO(wllenyj): Compatible with golang behavior.
if (msg.header.flags & FLAG_REMOTE_CLOSED) == FLAG_REMOTE_CLOSED
&& !msg.payload.is_empty()
@@ -518,7 +522,11 @@ impl HandlerContext {
}
}
- async fn handle_request(&self, msg: GenMessage) -> StdResult