Defer shutdown until actives sessions are 0

.github/workflows/ci.yml

@@ -31,7 +31,6 @@ concurrency:
 
 env:
   CARGO_TERM_COLOR: always
-  RUSTFLAGS: "-C target-feature=+aes,+avx2"
 
 jobs:
   lint:

build/build-image/Dockerfile

@@ -20,7 +20,6 @@ ENV RUSTUP_HOME=/usr/local/rustup \
     CARGO_HOME=/usr/local/cargo \
     PATH=/usr/local/cargo/bin:$PATH \
     CARGO_TERM_COLOR=always \
-    RUSTFLAGS="-C target-feature=+aes,+avx2" \
     LC_ALL=C.UTF-8 \
     LANG=C.UTF-8
 
@@ -80,7 +79,7 @@ RUN set -eux; \
     chmod -R a+w $RUSTUP_HOME $CARGO_HOME && \
     rustup component add rustfmt clippy && \
     rustup target add x86_64-unknown-linux-gnu aarch64-unknown-linux-gnu x86_64-pc-windows-gnu && \
-    cargo install mdbook cargo-watch && \
+    cargo install mdbook && \
     cargo install --locked cargo-about && \
     cargo install --locked mdbook-variables && \
     cargo install --locked cargo-deny && \

cloudbuild.yaml

@@ -99,7 +99,6 @@ options:
     - "REPOSITORY=${_REPOSITORY}"
     - "BUILD_IMAGE_TAG=${_BUILD_IMAGE_TAG}"
     - "CARGO_TERM_COLOR=always"
-    - 'RUSTFLAGS="-C target-feature=+aes,+avx2"'
   machineType: E2_HIGHCPU_32
   dynamic_substitutions: true
 timeout: 7200s

src/cli.rs

@@ -26,8 +26,13 @@ use crate::Config;
 use strum_macros::{Display, EnumString};
 
 pub use self::{
-    agent::Agent, generate_config_schema::GenerateConfigSchema, manage::Manage, proxy::Proxy,
-    qcmp::Qcmp, relay::Relay, service::Service,
+    agent::Agent,
+    generate_config_schema::GenerateConfigSchema,
+    manage::Manage,
+    proxy::Proxy,
+    qcmp::Qcmp,
+    relay::Relay,
+    service::{Finalizer, Service},
 };
 
 macro_rules! define_port {

src/cli/service.rs

@@ -115,6 +115,8 @@ pub struct Service {
     tls_key_path: Option<std::path::PathBuf>,
 }
 
+pub type Finalizer = Box<dyn FnOnce(&crate::signal::ShutdownRx) + Send>;
+
 impl Default for Service {
     fn default() -> Self {
         Self {
@@ -240,9 +242,9 @@ impl Service {
             config.id.store(id.into());
         }
 
-        let shutdown_rx = shutdown_rx.clone();
+        let mut shutdown_rx = shutdown_rx.clone();
         let mds_task = self.publish_mds(config)?;
-        let phoenix_task = self.publish_phoenix(config, &shutdown_rx)?;
+        let (phoenix_task, phoenix_finalizer) = self.publish_phoenix(config)?;
         // We need to call this before qcmp since if we use XDP we handle QCMP
         // internally without a separate task
         let (udp_task, finalizer) = self.publish_udp(config)?;
@@ -256,11 +258,18 @@ impl Service {
                 result = qcmp_task => result,
                 result = udp_task => result,
                 result = xds_task => result,
-            };
+                _ = shutdown_rx.changed() => {
+                    if let Some(finalizer) = finalizer {
+                        (finalizer)(&shutdown_rx);
 
-            if let Some(finalizer) = finalizer {
-                (finalizer)(shutdown_rx.clone());
-            }
+                        if let Some(pfin) = phoenix_finalizer {
+                            (pfin)(&shutdown_rx)
+                        }
+                    }
+
+                    Ok(())
+                }
+            };
 
             result
         }))
@@ -270,20 +279,23 @@ impl Service {
     fn publish_phoenix(
         &self,
         config: &Arc<Config>,
-        shutdown_rx: &crate::signal::ShutdownRx,
-    ) -> crate::Result<impl std::future::Future<Output = crate::Result<()>>> {
+    ) -> crate::Result<(
+        impl std::future::Future<Output = crate::Result<()>>,
+        Option<Finalizer>,
+    )> {
         if self.phoenix_enabled {
             tracing::info!(port=%self.qcmp_port, "starting phoenix service");
             let phoenix = crate::net::TcpListener::bind(Some(self.phoenix_port))?;
-            crate::net::phoenix::spawn(
+            let finalizer = crate::net::phoenix::spawn(
                 phoenix,
                 config.clone(),
-                shutdown_rx.clone(),
                 crate::net::phoenix::Phoenix::new(crate::codec::qcmp::QcmpMeasurement::new()?),
-            )?
+            )?;
+
+            return Ok((std::future::pending(), Some(finalizer)));
         }
 
-        Ok(std::future::pending())
+        Ok((std::future::pending(), None))
     }
 
     /// Spawns an QCMP server if enabled, otherwise returns a future which never completes.
@@ -357,10 +369,7 @@ impl Service {
     pub fn publish_udp(
         &mut self,
         config: &Arc<crate::config::Config>,
-    ) -> eyre::Result<(
-        impl Future<Output = crate::Result<()>>,
-        Option<Box<dyn FnOnce(crate::signal::ShutdownRx) + Send>>,
-    )> {
+    ) -> eyre::Result<(impl Future<Output = crate::Result<()>>, Option<Finalizer>)> {
         if !self.udp_enabled && !self.qcmp_enabled {
             return Ok((either::Left(std::future::pending()), None));
         }
@@ -403,10 +412,7 @@ impl Service {
     pub fn spawn_user_space_router(
         &self,
         config: Arc<crate::config::Config>,
-    ) -> crate::Result<(
-        impl Future<Output = crate::Result<()>>,
-        Box<dyn FnOnce(crate::signal::ShutdownRx) + Send>,
-    )> {
+    ) -> crate::Result<(impl Future<Output = crate::Result<()>>, Finalizer)> {
         let socket = crate::net::raw_socket_with_reuse(self.udp_port)?;
         let workers = self.udp_workers.get();
         let buffer_pool = Arc::new(crate::collections::BufferPool::new(workers, 2 * 1024));
@@ -431,7 +437,7 @@ impl Service {
 
         Ok((
             std::future::pending(),
-            Box::from(move |shutdown_rx: crate::signal::ShutdownRx| {
+            Box::from(move |shutdown_rx: &crate::signal::ShutdownRx| {
                 sessions.shutdown(*shutdown_rx.borrow() == crate::signal::ShutdownKind::Normal);
             }),
         ))
@@ -442,7 +448,7 @@ impl Service {
         &self,
         config: Arc<crate::config::Config>,
         force_xdp: bool,
-    ) -> eyre::Result<Box<dyn FnOnce(crate::signal::ShutdownRx) + Send>> {
+    ) -> eyre::Result<Finalizer> {
         use crate::net::xdp;
         use eyre::Context as _;
 
@@ -470,7 +476,7 @@ impl Service {
         .context("failed to setup XDP")?;
 
         let io_loop = xdp::spawn(workers, config).context("failed to spawn XDP I/O loop")?;
-        Ok(Box::new(move |srx: crate::signal::ShutdownRx| {
+        Ok(Box::new(move |srx: &crate::signal::ShutdownRx| {
             io_loop.shutdown(*srx.borrow() == crate::signal::ShutdownKind::Normal);
         }))
     }

src/components/proxy.rs

@@ -282,7 +282,7 @@ impl Proxy {
         )
         .port();
 
-        crate::cli::Service::default()
+        let svc_task = crate::cli::Service::default()
             .udp()
             .udp_port(udp_port)
             .xdp(self.xdp)
@@ -302,6 +302,10 @@ impl Proxy {
             .await
             .map_err(|error| eyre::eyre!(error))?;
 
+        if let Ok(Err(error)) = svc_task.await {
+            tracing::error!(%error, "Quilkin proxy services exited with error");
+        }
+
         Ok(())
     }
 }

src/components/proxy/packet_router.rs

@@ -73,6 +73,7 @@ impl<P: PacketMut> DownstreamPacket<P> {
         config: &Arc<Config>,
         sessions: &S,
         destinations: &mut Vec<crate::net::EndpointAddress>,
+        deny_new_sessions: bool,
     ) {
         tracing::trace!(
             id = worker_id,
@@ -82,7 +83,7 @@ impl<P: PacketMut> DownstreamPacket<P> {
         );
 
         let timer = metrics::processing_time(metrics::READ).start_timer();
-        if let Err(error) = self.process_inner(config, sessions, destinations) {
+        if let Err(error) = self.process_inner(config, sessions, destinations, deny_new_sessions) {
             let discriminant = error.discriminant();
 
             error.inc_system_errors_total(metrics::READ, &metrics::EMPTY);
@@ -99,6 +100,7 @@ impl<P: PacketMut> DownstreamPacket<P> {
         config: &Arc<Config>,
         sessions: &S,
         destinations: &mut Vec<crate::net::EndpointAddress>,
+        deny_new_sessions: bool,
     ) -> Result<(), PipelineError> {
         if !config.clusters.read().has_endpoints() {
             tracing::trace!("no upstream endpoints");
@@ -123,7 +125,7 @@ impl<P: PacketMut> DownstreamPacket<P> {
                 dest: epa.to_socket_addr()?,
             };
 
-            sessions.send(session_key, &contents)?;
+            sessions.send(session_key, &contents, deny_new_sessions)?;
         }
 
         Ok(())

src/components/proxy/packet_router/reference.rs

@@ -125,6 +125,7 @@ impl super::DownstreamReceiveWorkerConfig {
                                     &config,
                                     &sessions,
                                     &mut destinations,
+                                    true,
                                 );
                             }
                             Err(error) => {