@@ -9,6 +9,7 @@ use std::collections::HashMap;
99use std:: os:: unix:: io:: RawFd ;
1010use std:: result:: Result as StdResult ;
1111use std:: sync:: Arc ;
12+ use std:: time:: Duration ;
1213
1314use crate :: asynchronous:: stream:: { receive, respond, respond_with_status} ;
1415use crate :: asynchronous:: unix_incoming:: UnixIncoming ;
@@ -30,6 +31,7 @@ use tokio::{
3031 select, spawn,
3132 sync:: mpsc:: { channel, Receiver , Sender } ,
3233 sync:: watch,
34+ time:: timeout,
3335} ;
3436use tokio_vsock:: VsockListener ;
3537
@@ -322,10 +324,32 @@ async fn do_handle_request(
322324 metadata : context:: from_pb ( & req. metadata ) ,
323325 } ;
324326
325- method . handler ( ctx , req ) . await . map_err ( |e| {
327+ let get_unknown_status_and_log_err = |e| {
326328 error ! ( "method handle {} got error {:?}" , path, & e) ;
327329 get_status ( Code :: UNKNOWN , e)
328- } )
330+ } ;
331+
332+ if req. timeout_nano == 0 {
333+ method
334+ . handler ( ctx, req)
335+ . await
336+ . map_err ( get_unknown_status_and_log_err)
337+ } else {
338+ timeout (
339+ Duration :: from_nanos ( req. timeout_nano as u64 ) ,
340+ method. handler ( ctx, req) ,
341+ )
342+ . await
343+ . map_err ( |_| {
344+ // Timed out
345+ error ! ( "method handle {} got error timed out" , path) ;
346+ get_status ( Code :: DEADLINE_EXCEEDED , "timeout" )
347+ } )
348+ . and_then ( |r| {
349+ // Handler finished
350+ r. map_err ( get_unknown_status_and_log_err)
351+ } )
352+ }
329353}
330354
331355async fn handle_request (
0 commit comments