Merge pull request #205 from pwrdrvr/issue-185/improve-close-handling

Issue-185 - Extract close into a function

Author: huntharo

.github/workflows/build.yaml

@@ -47,6 +47,7 @@ jobs:
   rust-tests:
     needs: check-access
     runs-on: ubuntu-latest
+    timeout-minutes: 10
     env:
       CARGO_TERM_COLOR: always
     steps:

extension/src/lambda_request.rs

@@ -1,6 +1,7 @@
 use rand::Rng;
 use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
 use std::sync::Arc;
+use tokio::task::JoinHandle;
 
 use rand::SeedableRng;
 use tokio::{
@@ -12,8 +13,8 @@ use tokio_rustls::client::TlsStream;
 use crate::app_client::AppClient;
 use crate::endpoint::Endpoint;
 use crate::lambda_request_error::LambdaRequestError;
-use crate::messages;
-use crate::ping;
+use crate::messages::{self, ExitReason};
+use crate::ping::{self, send_close_request, PingResult};
 use crate::prelude::*;
 use crate::router_channel::RouterChannel;
 use crate::router_client::RouterClient;
@@ -43,6 +44,7 @@ pub struct LambdaRequest {
   requests_in_flight: Arc<AtomicUsize>,
   pub count: Arc<AtomicUsize>,
   start_time: u64,
+  last_active_grace_period_ms: u64,
 }
 
 impl LambdaRequest {
@@ -60,6 +62,7 @@ impl LambdaRequest {
     channel_count: u8,
     router_endpoint: Endpoint,
     deadline_ms: u64,
+    last_active_grace_period_ms: u64,
   ) -> Self {
     LambdaRequest {
       count: Arc::new(AtomicUsize::new(0)),
@@ -76,6 +79,7 @@ impl LambdaRequest {
       rng: rand::rngs::StdRng::from_entropy(),
       requests_in_flight: Arc::new(AtomicUsize::new(0)),
       start_time: current_time_millis(),
+      last_active_grace_period_ms,
     }
   }
 
@@ -86,7 +90,7 @@ impl LambdaRequest {
     router_client: RouterClient,
   ) -> Result<messages::ExitReason, LambdaRequestError> {
     // Send the ping requests in background
-    let ping_task = tokio::task::spawn(ping::send_ping_requests(
+    let mut ping_task = Some(tokio::task::spawn(ping::send_ping_requests(
       Arc::clone(&self.last_active),
       Arc::clone(&self.goaway_received),
       router_client.clone(),
@@ -97,7 +101,8 @@ impl LambdaRequest {
       self.deadline_ms,
       self.cancel_token.clone(),
       Arc::clone(&self.requests_in_flight),
-    ));
+      self.last_active_grace_period_ms,
+    )));
 
     // Startup the request channels
     let channel_futures = (0..self.channel_count)
@@ -169,7 +174,11 @@ impl LambdaRequest {
               if err.is_fatal() {
                 // Try to clean up the ping task
                 self.cancel_token.cancel();
-                let _ = tokio::time::timeout(std::time::Duration::from_secs(1), ping_task).await;
+                let _ = tokio::time::timeout(
+                  std::time::Duration::from_secs(1),
+                  ping_task.take().unwrap(),
+                )
+                .await;
 
                 return Err(err);
               }
@@ -189,12 +198,31 @@ impl LambdaRequest {
           err
         );
 
+        // If we got here we had a panic so we need to ask the router to
+        // shut down this Lambda invoke
+        let _ = send_close_request(
+          Arc::clone(&self.goaway_received),
+          router_client,
+          Arc::clone(&self.pool_id),
+          Arc::clone(&self.lambda_id),
+          self.router_endpoint.clone(),
+        )
+        .await;
+
         // Set the goaway signal so other tasks stop
         self.goaway_received.store(true, Ordering::Release);
 
         // Try to clean up the ping task
         self.cancel_token.cancel();
-        let _ = tokio::time::timeout(std::time::Duration::from_secs(1), ping_task).await;
+        let _ =
+          tokio::time::timeout(std::time::Duration::from_secs(1), ping_task.take().unwrap()).await;
+
+        // Shutdown the ping loop
+        // Note: we don't really care what happens here because
+        // we're already in a fatal error state
+        let _ = self
+          .wait_for_ping_loop(ping_task.take().unwrap(), exit_reason)
+          .await;
 
         // This is a bit of a lie: we don't know if the app connection is unreachable
         // but we do know we're in an non-deteministic state and that this
@@ -203,18 +231,36 @@ impl LambdaRequest {
       }
     }
 
-    // Wait for the ping loop to exit
+    // Shutdown the ping loop
+    exit_reason = self
+      .wait_for_ping_loop(ping_task.take().unwrap(), exit_reason)
+      .await
+      .unwrap_or(exit_reason)
+      .worse(exit_reason);
+
+    Ok(exit_reason)
+  }
+
+  pub async fn wait_for_ping_loop(
+    &mut self,
+    ping_task: JoinHandle<Option<PingResult>>,
+    exit_reason: ExitReason,
+  ) -> Result<ExitReason, LambdaRequestError> {
+    // Tell the ping loop to stop
     self.cancel_token.cancel();
+
+    // Wait for the ping loop to exit
     match ping_task.await {
       Ok(result) => {
         // Ping task completed successfully
 
         // If the ping task knows why we exited, use that reason
         if let Some(ping_result) = result {
           if let Some(ping_result) = ping_result.into() {
-            exit_reason = exit_reason.worse(ping_result);
+            return Ok(exit_reason.worse(ping_result));
           }
         }
+        Ok(exit_reason)
       }
       Err(e) => {
         log::error!(
@@ -223,11 +269,9 @@ impl LambdaRequest {
           e
         );
         // We'll lump this in as a generic router connection error
-        return Err(LambdaRequestError::RouterUnreachable.into());
+        Err(LambdaRequestError::RouterUnreachable)
       }
     }
-
-    Ok(exit_reason)
   }
 
   pub fn elapsed(&self) -> u64 {

extension/src/lambda_service.rs

@@ -188,6 +188,7 @@ impl LambdaService {
       channel_count,
       request.router_endpoint,
       deadline_ms,
+      self.options.last_active_grace_period_ms,
     );
 
     //
@@ -205,7 +206,7 @@ impl LambdaService {
           exit_reason
         );
 
-        resp.exit_reason = exit_reason;
+        resp.exit_reason = resp.exit_reason.worse(exit_reason);
       }
       Err(err) => {
         log::error!(

extension/src/options.rs

@@ -44,6 +44,7 @@ pub struct Options {
   pub local_env: bool,
   pub force_deadline_secs: Option<Duration>,
   pub async_init_timeout: Duration,
+  pub last_active_grace_period_ms: u64,
 }
 
 impl Options {
@@ -81,6 +82,11 @@ impl Options {
       local_env: provider.get_var("LAMBDA_DISPATCH_FORCE_DEADLINE").is_ok(),
       force_deadline_secs,
       async_init_timeout,
+      last_active_grace_period_ms: provider
+        .get_var("LAMBDA_DISPATCH_LAST_ACTIVE_GRACE_PERIOD_MS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .unwrap_or(250),
     }
   }
 }

extension/src/ping.rs

@@ -46,6 +46,62 @@ impl From<PingResult> for Option<messages::ExitReason> {
   }
 }
 
+pub async fn send_close_request(
+  goaway_received: Arc<AtomicBool>,
+  router_client: RouterClient,
+  pool_id: PoolId,
+  lambda_id: LambdaId,
+  router_endpoint: Endpoint,
+) {
+  let scheme = router_endpoint.scheme();
+  let host = router_endpoint.host();
+  let port = router_endpoint.port();
+  let host_header = router_endpoint.host_header();
+
+  let close_url = format!(
+    "{}://{}:{}/api/chunked/close/{}",
+    scheme.as_ref(),
+    host,
+    port,
+    lambda_id
+  );
+
+  // Send Close request to router
+  let (mut close_tx, close_recv) = mpsc::channel::<Result<Frame<Bytes>>>(1);
+  let boxed_close_body = BodyExt::boxed(StreamBody::new(close_recv));
+  let close_req = Request::builder()
+    .uri(&close_url)
+    .method("GET")
+    .header(hyper::header::DATE, fmt_http_date(SystemTime::now()))
+    .header("X-Pool-Id", pool_id.as_ref())
+    .header("X-Lambda-Id", lambda_id.as_ref());
+  let close_req = match &host_header {
+    Cow::Borrowed(v) => close_req.header(hyper::header::HOST, *v),
+    Cow::Owned(v) => close_req.header(hyper::header::HOST, v),
+  };
+  let close_req = close_req.body(boxed_close_body).unwrap();
+
+  let router_result = router_client.request(close_req).await;
+  close_tx.close().await.unwrap();
+  match router_result {
+    Ok(mut router_res) => {
+      // Rip through and discard so the response stream is closed
+      while router_res.frame().await.is_some() {}
+    }
+    Err(err) => {
+      log::error!(
+        "PoolId: {}, LambdaId: {} - PingLoop - Close request failed: {:?}",
+        pool_id,
+        lambda_id,
+        err
+      );
+    }
+  }
+
+  // Now mark that we are going away, after router has responded to our close request
+  goaway_received.store(true, Ordering::Release);
+}
+
 pub async fn send_ping_requests(
   last_active: Arc<AtomicU64>,
   goaway_received: Arc<AtomicBool>,
@@ -57,10 +113,11 @@ pub async fn send_ping_requests(
   deadline_ms: u64,
   cancel_token: tokio_util::sync::CancellationToken,
   requests_in_flight: Arc<AtomicUsize>,
+  last_active_grace_period_ms: u64,
 ) -> Option<PingResult> {
   let mut ping_result = None;
   let start_time = time::current_time_millis();
-  let mut last_ping_time = start_time;
+  let mut last_ping_time = 0;
 
   let scheme = router_endpoint.scheme();
   let host = router_endpoint.host();
@@ -76,17 +133,9 @@ pub async fn send_ping_requests(
     port,
     lambda_id
   );
-  let close_url = format!(
-    "{}://{}:{}/api/chunked/close/{}",
-    scheme.as_ref(),
-    host,
-    port,
-    lambda_id
-  );
 
   while !goaway_received.load(std::sync::atomic::Ordering::Acquire) && !cancel_token.is_cancelled()
   {
-    let last_active_grace_period_ms = 250;
     let close_before_deadline_ms = 15000;
     let last_active = last_active.load(Ordering::Acquire);
     let last_active_ago_ms = if last_active == 0 {
@@ -137,39 +186,14 @@ pub async fn send_ping_requests(
       }
 
       // Send Close request to router
-      let (mut close_tx, close_recv) = mpsc::channel::<Result<Frame<Bytes>>>(1);
-      let boxed_close_body = BodyExt::boxed(StreamBody::new(close_recv));
-      let close_req = Request::builder()
-        .uri(&close_url)
-        .method("GET")
-        .header(hyper::header::DATE, fmt_http_date(SystemTime::now()))
-        .header("X-Pool-Id", pool_id.as_ref())
-        .header("X-Lambda-Id", lambda_id.as_ref());
-      let close_req = match &host_header {
-        Cow::Borrowed(v) => close_req.header(hyper::header::HOST, *v),
-        Cow::Owned(v) => close_req.header(hyper::header::HOST, v),
-      };
-      let close_req = close_req.body(boxed_close_body).unwrap();
-
-      let router_result = router_client.request(close_req).await;
-      close_tx.close().await.unwrap();
-      match router_result {
-        Ok(mut router_res) => {
-          // Rip through and discard so the response stream is closed
-          while router_res.frame().await.is_some() {}
-        }
-        Err(err) => {
-          log::error!(
-            "PoolId: {}, LambdaId: {} - PingLoop - Close request failed: {:?}",
-            pool_id,
-            lambda_id,
-            err
-          );
-        }
-      }
-
-      // Now mark that we are going away, after router has responded to our close request
-      goaway_received.store(true, Ordering::Release);
+      send_close_request(
+        Arc::clone(&goaway_received),
+        router_client.clone(),
+        Arc::clone(&pool_id),
+        Arc::clone(&lambda_id),
+        router_endpoint.clone(),
+      )
+      .await;
       break;
     }
 
@@ -207,6 +231,8 @@ pub async fn send_ping_requests(
               lambda_id
             );
             ping_result.get_or_insert(PingResult::GoAway);
+
+            // This is from a goaway, so we do not need to ask the router to close our invoke
             goaway_received.store(true, Ordering::Release);
             break;
           }
@@ -225,6 +251,9 @@ pub async fn send_ping_requests(
               lambda_id,
               parts.status
             );
+
+            // TODO: This is not from a goaway, so we need to ask
+            // the router to close our invoke
             goaway_received.store(true, Ordering::Release);
             break;
           }
@@ -237,6 +266,9 @@ pub async fn send_ping_requests(
             lambda_id,
             err
           );
+
+          // TODO: This is not from a goaway, so we need to ask
+          // the router to close our invoke
           goaway_received.store(true, Ordering::Release);
           break;
         }