tests: update tests that presume a single AZ

Cargo.lock

@@ -7197,6 +7197,7 @@ dependencies = [
 [[package]]
 name = "tokio-epoll-uring"
 version = "0.1.0"
+source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#781989bb540a1408b0b93daa1e9d1fa452195497"
 dependencies = [
  "futures",
  "nix 0.26.4",
@@ -7807,6 +7808,7 @@ dependencies = [
 [[package]]
 name = "uring-common"
 version = "0.1.0"
+source = "git+https://github.com/neondatabase/tokio-epoll-uring.git?branch=main#781989bb540a1408b0b93daa1e9d1fa452195497"
 dependencies = [
  "bytes",
  "io-uring",

Cargo.toml

@@ -187,7 +187,7 @@ thiserror = "1.0"
 tikv-jemallocator = { version = "0.6", features = ["profiling", "stats", "unprefixed_malloc_on_supported_platforms"] }
 tikv-jemalloc-ctl = { version = "0.6", features = ["stats"] }
 tokio = { version = "1.43.1", features = ["macros"] }
-tokio-epoll-uring = { path = "../tokio-epoll-uring/tokio-epoll-uring" }
+tokio-epoll-uring = { git = "https://github.com/neondatabase/tokio-epoll-uring.git" , branch = "main" }
 tokio-io-timeout = "1.2.0"
 tokio-postgres-rustls = "0.12.0"
 tokio-rustls = { version = "0.26.0", default-features = false, features = ["tls12", "ring"]}

pageserver/src/metrics.rs

@@ -3867,8 +3867,6 @@ pub mod tokio_epoll_uring {
     use metrics::{Histogram, LocalHistogram, UIntGauge, register_histogram, register_int_counter};
     use once_cell::sync::Lazy;
 
-    use crate::virtual_file::io_engine::TokioEpollUringExtThreadLocalStateId;
-
     /// Shared storage for tokio-epoll-uring thread local metrics.
     pub(crate) static THREAD_LOCAL_METRICS_STORAGE: Lazy<ThreadLocalMetricsStorage> =
         Lazy::new(|| {
@@ -3888,7 +3886,7 @@ pub mod tokio_epoll_uring {
 
     pub struct ThreadLocalMetricsStorage {
         /// List of thread local metrics observers.
-        observers: Mutex<HashMap<TokioEpollUringExtThreadLocalStateId, Arc<ThreadLocalMetrics>>>,
+        observers: Mutex<HashMap<u64, Arc<ThreadLocalMetrics>>>,
         /// A histogram shared between all thread local systems
         /// for collecting slots submission queue depth.
         slots_submission_queue_depth: Histogram,
@@ -3910,10 +3908,7 @@ pub mod tokio_epoll_uring {
 
     impl ThreadLocalMetricsStorage {
         /// Registers a new thread local system. Returns a thread local metrics observer.
-        pub fn register_system(
-            &self,
-            id: TokioEpollUringExtThreadLocalStateId,
-        ) -> Arc<ThreadLocalMetrics> {
+        pub fn register_system(&self, id: u64) -> Arc<ThreadLocalMetrics> {
             let per_system_metrics = Arc::new(ThreadLocalMetrics::new(
                 self.slots_submission_queue_depth.local(),
             ));
@@ -3924,7 +3919,7 @@ pub mod tokio_epoll_uring {
 
         /// Removes metrics observer for a thread local system.
         /// This should be called before dropping a thread local system.
-        pub fn remove_system(&self, id: TokioEpollUringExtThreadLocalStateId) {
+        pub fn remove_system(&self, id: u64) {
             let mut g = self.observers.lock().unwrap();
             g.remove(&id);
         }

pageserver/src/virtual_file/io_engine.rs

@@ -11,8 +11,6 @@
 
 #[cfg(target_os = "linux")]
 pub(super) mod tokio_epoll_uring_ext;
-#[cfg(target_os = "linux")]
-pub(crate) use tokio_epoll_uring_ext::ThreadLocalStateId as TokioEpollUringExtThreadLocalStateId;
 
 use tokio_epoll_uring::IoBuf;
 use tracing::Instrument;

pageserver/src/virtual_file/io_engine/tokio_epoll_uring_ext.rs

@@ -5,9 +5,8 @@
 //! on older kernels, such as some (but not all) older kernels in the Linux 5.10 series.
 //! See <https://github.com/neondatabase/neon/issues/6373#issuecomment-1905814391> for more details.
 
-use std::collections::HashMap;
+use std::sync::Arc;
 use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
-use std::sync::{Arc, Weak};
 
 use tokio_epoll_uring::{System, SystemHandle};
 use tokio_util::sync::CancellationToken;
@@ -23,14 +22,13 @@ struct ThreadLocalState(Arc<ThreadLocalStateInner>);
 struct ThreadLocalStateInner {
     cell: tokio::sync::OnceCell<SystemHandle<metrics::ThreadLocalMetrics>>,
     launch_attempts: AtomicU32,
-    thread_local_state_id: ThreadLocalStateId,
+    /// populated through fetch_add from [`THREAD_LOCAL_STATE_ID`]
+    thread_local_state_id: u64,
 }
 
 impl Drop for ThreadLocalStateInner {
     fn drop(&mut self) {
         THREAD_LOCAL_METRICS_STORAGE.remove_system(self.thread_local_state_id);
-        let mut launched_systems = ALL_THREAD_LOCALS.lock().unwrap();
-        launched_systems.remove(&self.thread_local_state_id);
     }
 }
 
@@ -39,7 +37,7 @@ impl ThreadLocalState {
         Self(Arc::new(ThreadLocalStateInner {
             cell: tokio::sync::OnceCell::default(),
             launch_attempts: AtomicU32::new(0),
-            thread_local_state_id: ThreadLocalStateId::new(),
+            thread_local_state_id: THREAD_LOCAL_STATE_ID.fetch_add(1, Ordering::Relaxed),
         }))
     }
 
@@ -48,23 +46,7 @@ impl ThreadLocalState {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct ThreadLocalStateId(u64);
-impl ThreadLocalStateId {
-    pub fn new() -> Self {
-        static THREAD_LOCAL_STATE_ID: AtomicU64 = AtomicU64::new(0);
-        Self(THREAD_LOCAL_STATE_ID.fetch_add(1, Ordering::Relaxed))
-    }
-}
-impl std::fmt::Display for ThreadLocalStateId {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}", self.0)
-    }
-}
-
-static ALL_THREAD_LOCALS: std::sync::LazyLock<
-    std::sync::Mutex<HashMap<ThreadLocalStateId, Weak<ThreadLocalStateInner>>>,
-> = std::sync::LazyLock::new(|| Default::default());
+static THREAD_LOCAL_STATE_ID: AtomicU64 = AtomicU64::new(0);
 
 thread_local! {
     static THREAD_LOCAL: ThreadLocalState = ThreadLocalState::new();
@@ -94,23 +76,7 @@ pub async fn thread_local_system() -> Handle {
                     )
                     .await;
                     let per_system_metrics = metrics::THREAD_LOCAL_METRICS_STORAGE.register_system(inner.thread_local_state_id);
-                    let res = System::launch_with_metrics(per_system_metrics,
-                        if utils::env::var("NEON_PAGESERVER_TOKIO_EPOLL_URING_SHARE_IO_WQ_POOL").unwrap_or(false) {
-                        ALL_THREAD_LOCALS.lock().unwrap().values().filter_map(|other: &Weak<ThreadLocalStateInner>| {
-                            let Some(other) = other.upgrade() else {
-                                return None;
-                            };
-                            if other.thread_local_state_id == inner.thread_local_state_id {
-                                return None;
-                            }
-                            let Some(system) = inner.cell.get() else {
-                                return None;
-                            };
-                            Some(system)
-                        }).next()
-                    } else {
-                        None
-                    },)
+                    let res = System::launch_with_metrics(per_system_metrics)
                     // this might move us to another executor thread => loop outside the get_or_try_init, not inside it
                     .await;
                     match res {

test_runner/fixtures/neon_fixtures.py

@@ -435,7 +435,7 @@ class NeonEnvBuilder:
         pageserver_config_override: str | Callable[[dict[str, Any]], None] | None = None,
         num_safekeepers: int = 1,
         num_pageservers: int = 1,
-        num_azs: int = 1,
+        num_azs: int = 3,
         # Use non-standard SK ids to check for various parsing bugs
         safekeepers_id_start: int = 0,
         # fsync is disabled by default to make the tests go faster

test_runner/regress/test_sharding.py

@@ -44,6 +44,10 @@ def test_sharding_smoke(
     """
 
     shard_count = 4
+
+    # Single AZ so that a sharded tenant can spread across all pageservers
+    neon_env_builder.num_azs = 1
+
     neon_env_builder.num_pageservers = shard_count
 
     # 1MiB stripes: enable getting some meaningful data distribution without
@@ -523,16 +527,13 @@ def test_sharding_split_smoke(
     # In preferred AZ & other AZ we will end up with one shard per pageserver
     neon_env_builder.num_pageservers = split_shard_count * 2
 
-    # Two AZs
-    def assign_az(ps_cfg):
-        az = f"az-{(ps_cfg['id'] - 1) % 2}"
-        ps_cfg["availability_zone"] = az
-
+    def ps_conf_hook(ps_cfg):
         # We will run more pageservers than tests usually do, so give them tiny page caches
         # in case we're on a test node under memory pressure.
         ps_cfg["page_cache_size"] = 128
 
-    neon_env_builder.pageserver_config_override = assign_az
+    neon_env_builder.num_azs = 2
+    neon_env_builder.pageserver_config_override = ps_conf_hook
 
     # 1MiB stripes: enable getting some meaningful data distribution without
     # writing large quantities of data in this test.  The stripe size is given

test_runner/regress/test_storage_controller.py

@@ -2212,6 +2212,7 @@ def test_graceful_cluster_restart(
     tenants where we fill based on a target shard count.
     """
     neon_env_builder.num_azs = num_azs
+
     neon_env_builder.num_pageservers = 2
     neon_env_builder.control_plane_hooks_api = compute_reconfigure_listener.control_plane_hooks_api
     env = neon_env_builder.init_configs()
@@ -2475,7 +2476,6 @@ def test_storage_controller_node_deletion(
     Test that deleting a node works & properly reschedules everything that was on the node.
     """
     neon_env_builder.num_pageservers = 3
-    neon_env_builder.num_azs = 3
     env = neon_env_builder.init_configs()
     env.start()
 
@@ -3512,12 +3512,7 @@ def eq_safekeeper_records(a: dict[str, Any], b: dict[str, Any]) -> bool:
 
 @run_only_on_default_postgres("this is like a 'unit test' against storcon db")
 def test_shard_preferred_azs(neon_env_builder: NeonEnvBuilder):
-    def assign_az(ps_cfg):
-        az = f"az-{ps_cfg['id'] % 2}"
-        log.info("Assigned AZ {az}")
-        ps_cfg["availability_zone"] = az
-
-    neon_env_builder.pageserver_config_override = assign_az
+    neon_env_builder.num_azs = 2
     neon_env_builder.num_pageservers = 4
     env = neon_env_builder.init_configs()
     env.start()

test_runner/regress/test_timeline_detach_ancestor.py

@@ -1186,6 +1186,8 @@ def test_sharded_tad_interleaved_after_partial_success(neon_env_builder: NeonEnv
     """
 
     shard_count = 2
+    # Single AZ so that a sharded tenant can spread across all pageservers
+    neon_env_builder.num_azs = 1
     neon_env_builder.num_pageservers = shard_count
     env = neon_env_builder.init_start(initial_tenant_shard_count=shard_count)