Add shutdown flag to proxy session

Add a Session::is_process_shutting_down API for custom logic to be
actioned in filters when shutdown has started.

Author: drcaramelsyrup

.bleep

@@ -1 +1 @@
-93b912130db516aecdfb9b87b48fa7797f9d1a49
+03479bcb7e6312754af45a43f6cdd173538a3cc2

pingora-proxy/src/lib.rs

@@ -45,7 +45,10 @@ use once_cell::sync::Lazy;
 use pingora_http::{RequestHeader, ResponseHeader};
 use std::fmt::Debug;
 use std::str;
-use std::sync::Arc;
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc,
+};
 use std::time::Duration;
 use tokio::sync::{mpsc, Notify};
 use tokio::time;
@@ -111,6 +114,7 @@ where
     inner: SV, // TODO: name it better than inner
     client_upstream: Connector<C>,
     shutdown: Notify,
+    shutdown_flag: Arc<AtomicBool>,
     pub server_options: Option<HttpServerOptions>,
     pub h2_options: Option<H2Options>,
     pub downstream_modules: HttpModules,
@@ -124,6 +128,7 @@ impl<SV> HttpProxy<SV, ()> {
             inner,
             client_upstream: Connector::new(Some(ConnectorOptions::from_server_conf(&conf))),
             shutdown: Notify::new(),
+            shutdown_flag: Arc::new(AtomicBool::new(false)),
             server_options: None,
             h2_options: None,
             downstream_modules: HttpModules::new(),
@@ -154,6 +159,7 @@ where
             inner,
             client_upstream,
             shutdown: Notify::new(),
+            shutdown_flag: Arc::new(AtomicBool::new(false)),
             server_options: None,
             downstream_modules: HttpModules::new(),
             max_retries: conf.max_retries,
@@ -401,12 +407,15 @@ pub struct Session {
     /// Upstream response body bytes received (payload only). Set by proxy layer.
     /// TODO: move this into an upstream session digest for future fields.
     upstream_body_bytes_received: usize,
+    /// Flag that is set when the shutdown process has begun.
+    shutdown_flag: Arc<AtomicBool>,
 }
 
 impl Session {
     fn new(
         downstream_session: impl Into<Box<HttpSession>>,
         downstream_modules: &HttpModules,
+        shutdown_flag: Arc<AtomicBool>,
     ) -> Self {
         Session {
             downstream_session: downstream_session.into(),
@@ -419,22 +428,33 @@ impl Session {
             subrequest_spawner: None, // optionally set later on
             downstream_modules_ctx: downstream_modules.build_ctx(),
             upstream_body_bytes_received: 0,
+            shutdown_flag,
         }
     }
 
     /// Create a new [Session] from the given [Stream]
     ///
-    /// This function is mostly used for testing and mocking.
+    /// This function is mostly used for testing and mocking, given the downstream modules and
+    /// shutdown flags will never be set.
     pub fn new_h1(stream: Stream) -> Self {
         let modules = HttpModules::new();
-        Self::new(Box::new(HttpSession::new_http1(stream)), &modules)
+        Self::new(
+            Box::new(HttpSession::new_http1(stream)),
+            &modules,
+            Arc::new(AtomicBool::new(false)),
+        )
     }
 
     /// Create a new [Session] from the given [Stream] with modules
     ///
-    /// This function is mostly used for testing and mocking.
+    /// This function is mostly used for testing and mocking, given the shutdown flag will never be
+    /// set.
     pub fn new_h1_with_modules(stream: Stream, downstream_modules: &HttpModules) -> Self {
-        Self::new(Box::new(HttpSession::new_http1(stream)), downstream_modules)
+        Self::new(
+            Box::new(HttpSession::new_http1(stream)),
+            downstream_modules,
+            Arc::new(AtomicBool::new(false)),
+        )
     }
 
     pub fn as_downstream_mut(&mut self) -> &mut HttpSession {
@@ -569,6 +589,11 @@ impl Session {
         self.upstream_body_bytes_received = n;
     }
 
+    /// Is the proxy process in the process of shutting down (e.g. due to graceful upgrade)?
+    pub fn is_process_shutting_down(&self) -> bool {
+        self.shutdown_flag.load(Ordering::Acquire)
+    }
+
     pub fn downstream_custom_message(
         &mut self,
     ) -> Result<
@@ -910,7 +935,11 @@ where
         debug!("starting subrequest");
 
         let mut session = match self.handle_new_request(session).await {
-            Some(downstream_session) => Session::new(downstream_session, &self.downstream_modules),
+            Some(downstream_session) => Session::new(
+                downstream_session,
+                &self.downstream_modules,
+                self.shutdown_flag.clone(),
+            ),
             None => return, // bad request
         };
 
@@ -978,7 +1007,11 @@ where
 
         // TODO: keepalive pool, use stack
         let mut session = match self.handle_new_request(session).await {
-            Some(downstream_session) => Session::new(downstream_session, &self.downstream_modules),
+            Some(downstream_session) => Session::new(
+                downstream_session,
+                &self.downstream_modules,
+                self.shutdown_flag.clone(),
+            ),
             None => return None, // bad request
         };
 
@@ -992,10 +1025,9 @@ where
     }
 
     async fn http_cleanup(&self) {
+        self.shutdown_flag.store(true, Ordering::Release);
         // Notify all keepalived requests blocking on read_request() to abort
         self.shutdown.notify_waiters();
-
-        // TODO: impl shutting down flag so that we don't need to read stack.is_shutting_down()
     }
 
     fn server_options(&self) -> Option<&HttpServerOptions> {