performance: reduce and test size of futures (istio#985)

* performance: reduce and test size of futures

Prior to this PR, each outbound connection future was 22k. This is
really big.

The reason behind this is how Rust sizes Futures; the stack size is
allocated based on the maximum required size through any branch. This
means even obscure code paths can bloat up the entire future for
everyone.

This PR optimizes a bunch of these. This is done in a few ways:
* Box::pin expensive futures. While we have to allocate these when we
  hit them, of course, when we don't hit these codepaths they are now
"free". We have a lot of uncommon expensive code paths. Especially HBONE
vs passthrough; hbone is way more expensive. But also on demand XDS,
DNS, etc
* Arc the Config which is large
* A variety of other micro-optimizations

A new `assertions` module is added to verify we do not regress

* Fix the upgrade drop issue

* fmt

Cargo.toml

@@ -57,7 +57,7 @@ http-body-04 = { package = "http-body", version = "0.4" }
 http-body-1 = { package = "http-body", version = "1.0.0-rc.2" }
 http-body-util = "0.1"
 http-types = { version = "2.12", default-features = false }
-hyper = { version = "1.2", features = ["full"] }
+hyper = { version = "1.3", features = ["full"] }
 hyper-rustls = { version = "0.27.0", default-features = false, features = ["logging", "http1", "http2"] }
 hyper-util = { version = "0.1", features = ["full"] }
 ipnet = { version = "2.9", features = ["serde"] }

benches/throughput.rs

@@ -30,7 +30,6 @@ use tokio::runtime::Runtime;
 use tokio::sync::Mutex;
 use tracing::info;
 
-use ztunnel::metrics::IncrementRecorder;
 use ztunnel::rbac::{Authorization, RbacMatch, StringMatch};
 use ztunnel::state::workload::{Protocol, Workload};
 use ztunnel::test_helpers::app::TestApp;
@@ -127,7 +126,7 @@ fn initialize_environment(
             None,
             config_source,
         );
-        let app = app::build_with_cert(config, cert_manager.clone())
+        let app = app::build_with_cert(Arc::new(config), cert_manager.clone())
             .await
             .unwrap();
 
@@ -363,14 +362,16 @@ pub fn metrics(c: &mut Criterion) {
     let mut c = c.benchmark_group("metrics");
     c.bench_function("write", |b| {
         b.iter(|| {
-            metrics.increment(&proxy::ConnectionOpen {
+            let co = proxy::ConnectionOpen {
                 reporter: Default::default(),
                 source: Some(test_helpers::test_default_workload()),
                 derived_source: None,
                 destination: None,
                 destination_service: None,
                 connection_security_policy: Default::default(),
-            })
+            };
+            let tl = proxy::CommonTrafficLabels::from(co);
+            metrics.connection_opens.get_or_create(&tl).inc();
         })
     });
     c.bench_function("encode", |b| {
@@ -455,7 +456,7 @@ fn hbone_connections(c: &mut Criterion) {
             None,
             config_source,
         );
-        let app = app::build_with_cert(config, cert_manager.clone())
+        let app = app::build_with_cert(Arc::new(config), cert_manager.clone())
             .await
             .unwrap();
         let ta = TestApp::from((&app, cert_manager));

src/admin.rs

@@ -62,7 +62,7 @@ pub trait AdminHandler2: Sync + Send {
 
 struct State {
     proxy_state: DemandProxyState,
-    config: Config,
+    config: Arc<Config>,
     shutdown_trigger: signal::ShutdownTrigger,
     cert_manager: Arc<SecretManager>,
     handlers: Vec<Arc<dyn AdminHandler2>>,
@@ -79,7 +79,7 @@ pub struct ConfigDump {
     proxy_state: DemandProxyState,
     static_config: LocalConfig,
     version: BuildInfo,
-    config: Config,
+    config: Arc<Config>,
     certificates: Vec<CertsDump>,
 }
 
@@ -103,7 +103,7 @@ pub struct CertsDump {
 
 impl Service {
     pub async fn new(
-        config: Config,
+        config: Arc<Config>,
         proxy_state: DemandProxyState,
         shutdown_trigger: signal::ShutdownTrigger,
         drain_rx: Watch,
@@ -478,6 +478,7 @@ mod tests {
     use bytes::Bytes;
     use http_body_util::BodyExt;
     use std::collections::HashMap;
+    use std::sync::Arc;
     use std::time::Duration;
 
     fn diff_json<'a>(a: &'a serde_json::Value, b: &'a serde_json::Value) -> String {
@@ -744,7 +745,7 @@ mod tests {
             proxy_state,
             static_config: Default::default(),
             version: Default::default(),
-            config: default_config,
+            config: Arc::new(default_config),
             certificates: dump_certs(&manager).await,
         };
 

src/app.rs

@@ -33,7 +33,7 @@ use crate::{admin, config, metrics, proxy, readiness, signal};
 use crate::{dns, xds};
 
 pub async fn build_with_cert(
-    config: config::Config,
+    config: Arc<config::Config>,
     cert_manager: Arc<SecretManager>,
 ) -> anyhow::Result<Bound> {
     // Start the data plane worker pool.
@@ -280,7 +280,7 @@ fn new_data_plane_pool(num_worker_threads: usize) -> mpsc::Sender<DataPlaneTask>
     tx
 }
 
-pub async fn build(config: config::Config) -> anyhow::Result<Bound> {
+pub async fn build(config: Arc<config::Config>) -> anyhow::Result<Bound> {
     let cert_manager = if config.fake_ca {
         mock_secret_manager()
     } else {

src/assertions.rs

@@ -0,0 +1,47 @@
+// Copyright Istio Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Below helper functions are used to help make about the size of types.
+// There are some compile time ways to do this, but they don't work in the way we need for the most part;
+// analyzing the size of Futures which we don't have explicit declarations for.
+// Future size is determined by the max required stack size for the async function. This means deeply
+// branched code can create huge Future's, leading to high per-connection memory usage in ztunnel.
+// Debugging these usages can be done by `RUSTFLAGS=-Zprint-type-sizes cargo +nightly build -j 1`,
+// or by logging with the functions below.
+
+#[cfg(all(any(test, feature = "testing"), debug_assertions))]
+pub fn size_between_ref<T>(min: usize, max: usize, t: &T) {
+    let size = std::mem::size_of_val(t);
+    if size < min || size > max {
+        // If it is too small: that is good, we just want to update the assertion to be more aggressive
+        // If it is too big: that is bad. We may need to increase the limit, or consider refactors.
+        panic!(
+            "type {} size is unexpected, wanted {min}..{max}, got {size}",
+            std::any::type_name::<T>(),
+        )
+    }
+    tracing::trace!(
+        "type {} size is within expectations, wanted {min}..{max}, got {size}",
+        std::any::type_name::<T>(),
+    )
+}
+
+#[cfg(not(all(any(test, feature = "testing"), debug_assertions)))]
+pub fn size_between_ref<T>(_min: usize, _max: usize, _t: &T) {}
+
+#[inline(always)]
+pub fn size_between<T>(min: usize, max: usize, t: T) -> T {
+    size_between_ref(min, max, &t);
+    t
+}

src/identity/manager.rs

@@ -465,13 +465,15 @@ impl fmt::Debug for SecretManager {
 }
 
 impl SecretManager {
-    pub async fn new(cfg: crate::config::Config) -> Result<Self, Error> {
+    pub async fn new(cfg: Arc<crate::config::Config>) -> Result<Self, Error> {
         let caclient = CaClient::new(
-            cfg.ca_address.expect("ca_address must be set to use CA"),
+            cfg.ca_address
+                .clone()
+                .expect("ca_address must be set to use CA"),
             Box::new(tls::ControlPlaneAuthentication::RootCert(
                 cfg.ca_root_cert.clone(),
             )),
-            cfg.auth,
+            cfg.auth.clone(),
             cfg.proxy_mode == ProxyMode::Shared,
             cfg.secret_ttl.as_secs().try_into().unwrap_or(60 * 60 * 24),
         )

src/inpod/test_helpers.rs

@@ -76,7 +76,7 @@ impl Default for Fixture {
 
         let ipc = InPodConfig::new(&cfg).unwrap();
         let proxy_gen = ProxyFactory::new(
-            cfg,
+            Arc::new(cfg),
             dstate,
             cert_manager,
             Some(metrics),

src/lib.rs

@@ -14,6 +14,7 @@
 
 pub mod admin;
 pub mod app;
+pub mod assertions;
 pub mod baggage;
 pub mod cert_fetcher;
 pub mod config;

src/main.rs

@@ -14,11 +14,11 @@
 
 extern crate core;
 
+use std::sync::Arc;
 use tracing::info;
 use ztunnel::*;
 
 #[cfg(feature = "jemalloc")]
-use tikv_jemallocator;
 #[cfg(feature = "jemalloc")]
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
@@ -30,7 +30,7 @@ pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0
 
 fn main() -> anyhow::Result<()> {
     telemetry::setup_logging();
-    let config: config::Config = config::parse_config()?;
+    let config = Arc::new(config::parse_config()?);
 
     // For now we don't need a complex CLI, so rather than pull in dependencies just use basic argv[1]
     match std::env::args().nth(1).as_deref() {
@@ -70,7 +70,7 @@ fn version() -> anyhow::Result<()> {
     Ok(())
 }
 
-async fn proxy(cfg: config::Config) -> anyhow::Result<()> {
+async fn proxy(cfg: Arc<config::Config>) -> anyhow::Result<()> {
     info!("version: {}", version::BuildInfo::new());
     info!("running with config: {}", serde_yaml::to_string(&cfg)?);
     app::build(cfg).await?.wait_termination().await

src/metrics/server.rs

@@ -31,7 +31,11 @@ pub struct Server {
 }
 
 impl Server {
-    pub async fn new(config: Config, drain_rx: Watch, registry: Registry) -> anyhow::Result<Self> {
+    pub async fn new(
+        config: Arc<Config>,
+        drain_rx: Watch,
+        registry: Registry,
+    ) -> anyhow::Result<Self> {
         hyper_util::Server::<Mutex<Registry>>::bind(
             "stats",
             config.stats_addr,

src/proxy.rs

@@ -99,7 +99,7 @@ pub struct Proxy {
 
 #[derive(Clone)]
 pub(super) struct ProxyInputs {
-    cfg: config::Config,
+    cfg: Arc<config::Config>,
     cert_manager: Arc<SecretManager>,
     connection_manager: ConnectionManager,
     hbone_port: u16,
@@ -112,7 +112,7 @@ pub(super) struct ProxyInputs {
 #[allow(clippy::too_many_arguments)]
 impl ProxyInputs {
     pub fn new(
-        cfg: config::Config,
+        cfg: Arc<config::Config>,
         cert_manager: Arc<SecretManager>,
         connection_manager: ConnectionManager,
         state: DemandProxyState,
@@ -135,7 +135,7 @@ impl ProxyInputs {
 
 impl Proxy {
     pub async fn new(
-        cfg: config::Config,
+        cfg: Arc<config::Config>,
         state: DemandProxyState,
         cert_manager: Arc<SecretManager>,
         metrics: Metrics,