From f85935446690d30b4874ed2cfa76fd15b187a5e8 Mon Sep 17 00:00:00 2001 From: "Alex Chi Z." <4198311+skyzh@users.noreply.github.com> Date: Tue, 22 Jul 2025 13:55:07 -0400 Subject: [PATCH] feat(pageserver): add db rel count as feature flag property (#12632) ## Problem As part of the reldirv2 rollout: LKB-197. We will use number of db/rels as a criteria whether to rollout reldirv2 directly on the write path (simplest and easiest way of rollout). If the number of rel/db is small then it shouldn't take too long time on the write path. ## Summary of changes * Compute db/rel count during basebackup. * Also compute it during logical size computation. * Collect maximum number of db/rel across all timelines in the feature flag propeties. --------- Signed-off-by: Alex Chi Z --- pageserver/src/basebackup.rs | 12 +++++++- pageserver/src/feature_resolver.rs | 38 +++++++++++++++++++++++- pageserver/src/pgdatadir_mapping.rs | 10 +++++++ pageserver/src/tenant/timeline.rs | 16 +++++++--- test_runner/regress/test_feature_flag.py | 4 +++ 5 files changed, 74 insertions(+), 6 deletions(-) diff --git a/pageserver/src/basebackup.rs b/pageserver/src/basebackup.rs index 1a44c80e2d..1f1a3f8157 100644 --- a/pageserver/src/basebackup.rs +++ b/pageserver/src/basebackup.rs @@ -11,6 +11,7 @@ //! from data stored in object storage. //! use std::fmt::Write as FmtWrite; +use std::sync::Arc; use std::time::{Instant, SystemTime}; use anyhow::{Context, anyhow}; @@ -420,12 +421,16 @@ where } let mut min_restart_lsn: Lsn = Lsn::MAX; + + let mut dbdir_cnt = 0; + let mut rel_cnt = 0; + // Create tablespace directories for ((spcnode, dbnode), has_relmap_file) in self.timeline.list_dbdirs(self.lsn, self.ctx).await? { self.add_dbdir(spcnode, dbnode, has_relmap_file).await?; - + dbdir_cnt += 1; // If full backup is requested, include all relation files. // Otherwise only include init forks of unlogged relations. let rels = self @@ -433,6 +438,7 @@ where .list_rels(spcnode, dbnode, Version::at(self.lsn), self.ctx) .await?; for &rel in rels.iter() { + rel_cnt += 1; // Send init fork as main fork to provide well formed empty // contents of UNLOGGED relations. Postgres copies it in // `reinit.c` during recovery. @@ -455,6 +461,10 @@ where } } + self.timeline + .db_rel_count + .store(Some(Arc::new((dbdir_cnt, rel_cnt)))); + let start_time = Instant::now(); let aux_files = self .timeline diff --git a/pageserver/src/feature_resolver.rs b/pageserver/src/feature_resolver.rs index 11b0e972b4..678d7e052b 100644 --- a/pageserver/src/feature_resolver.rs +++ b/pageserver/src/feature_resolver.rs @@ -156,6 +156,8 @@ impl FeatureResolver { let tenant_properties = PerTenantProperties { remote_size_mb: Some(rand::rng().random_range(100.0..1000000.00)), + db_count_max: Some(rand::rng().random_range(1..1000)), + rel_count_max: Some(rand::rng().random_range(1..1000)), } .into_posthog_properties(); @@ -344,6 +346,8 @@ impl FeatureResolver { struct PerTenantProperties { pub remote_size_mb: Option, + pub db_count_max: Option, + pub rel_count_max: Option, } impl PerTenantProperties { @@ -355,6 +359,18 @@ impl PerTenantProperties { PostHogFlagFilterPropertyValue::Number(remote_size_mb), ); } + if let Some(db_count) = self.db_count_max { + properties.insert( + "tenant_db_count_max".to_string(), + PostHogFlagFilterPropertyValue::Number(db_count as f64), + ); + } + if let Some(rel_count) = self.rel_count_max { + properties.insert( + "tenant_rel_count_max".to_string(), + PostHogFlagFilterPropertyValue::Number(rel_count as f64), + ); + } properties } } @@ -409,7 +425,11 @@ impl TenantFeatureResolver { /// Refresh the cached properties and flags on the critical path. pub fn refresh_properties_and_flags(&self, tenant_shard: &TenantShard) { + // Any of the remote size is none => this property is none. let mut remote_size_mb = Some(0.0); + // Any of the db or rel count is available => this property is available. + let mut db_count_max = None; + let mut rel_count_max = None; for timeline in tenant_shard.list_timelines() { let size = timeline.metrics.resident_physical_size_get(); if size == 0 { @@ -419,9 +439,25 @@ impl TenantFeatureResolver { if let Some(ref mut remote_size_mb) = remote_size_mb { *remote_size_mb += size as f64 / 1024.0 / 1024.0; } + if let Some(data) = timeline.db_rel_count.load_full() { + let (db_count, rel_count) = *data.as_ref(); + if db_count_max.is_none() { + db_count_max = Some(db_count); + } + if rel_count_max.is_none() { + rel_count_max = Some(rel_count); + } + db_count_max = db_count_max.map(|max| max.max(db_count)); + rel_count_max = rel_count_max.map(|max| max.max(rel_count)); + } } self.cached_tenant_properties.store(Arc::new( - PerTenantProperties { remote_size_mb }.into_posthog_properties(), + PerTenantProperties { + remote_size_mb, + db_count_max, + rel_count_max, + } + .into_posthog_properties(), )); // BEGIN: Update the feature flag on the critical path. diff --git a/pageserver/src/pgdatadir_mapping.rs b/pageserver/src/pgdatadir_mapping.rs index ab9cc88e5f..c9f3184188 100644 --- a/pageserver/src/pgdatadir_mapping.rs +++ b/pageserver/src/pgdatadir_mapping.rs @@ -8,6 +8,7 @@ //! use std::collections::{HashMap, HashSet, hash_map}; use std::ops::{ControlFlow, Range}; +use std::sync::Arc; use crate::walingest::{WalIngestError, WalIngestErrorKind}; use crate::{PERF_TRACE_TARGET, ensure_walingest}; @@ -1254,11 +1255,16 @@ impl Timeline { let dbdir = DbDirectory::des(&buf)?; let mut total_size: u64 = 0; + let mut dbdir_cnt = 0; + let mut rel_cnt = 0; + for (spcnode, dbnode) in dbdir.dbdirs.keys() { + dbdir_cnt += 1; for rel in self .list_rels(*spcnode, *dbnode, Version::at(lsn), ctx) .await? { + rel_cnt += 1; if self.cancel.is_cancelled() { return Err(CalculateLogicalSizeError::Cancelled); } @@ -1269,6 +1275,10 @@ impl Timeline { total_size += relsize as u64; } } + + self.db_rel_count + .store(Some(Arc::new((dbdir_cnt, rel_cnt)))); + Ok(total_size * BLCKSZ as u64) } diff --git a/pageserver/src/tenant/timeline.rs b/pageserver/src/tenant/timeline.rs index 7f6173db3f..3ef07aa414 100644 --- a/pageserver/src/tenant/timeline.rs +++ b/pageserver/src/tenant/timeline.rs @@ -287,7 +287,7 @@ pub struct Timeline { ancestor_lsn: Lsn, // The LSN of gc-compaction that was last applied to this timeline. - gc_compaction_state: ArcSwap>, + gc_compaction_state: ArcSwapOption, pub(crate) metrics: Arc, @@ -450,6 +450,9 @@ pub struct Timeline { #[expect(dead_code)] feature_resolver: Arc, + + /// Basebackup will collect the count and store it here. Used for reldirv2 rollout. + pub(crate) db_rel_count: ArcSwapOption<(usize, usize)>, } pub(crate) enum PreviousHeatmap { @@ -3237,7 +3240,7 @@ impl Timeline { }), disk_consistent_lsn: AtomicLsn::new(disk_consistent_lsn.0), - gc_compaction_state: ArcSwap::new(Arc::new(gc_compaction_state)), + gc_compaction_state: ArcSwapOption::from_pointee(gc_compaction_state), last_freeze_at: AtomicLsn::new(disk_consistent_lsn.0), last_freeze_ts: RwLock::new(Instant::now()), @@ -3342,6 +3345,8 @@ impl Timeline { basebackup_cache: resources.basebackup_cache, feature_resolver: resources.feature_resolver.clone(), + + db_rel_count: ArcSwapOption::from_pointee(None), }; result.repartition_threshold = @@ -3413,7 +3418,7 @@ impl Timeline { gc_compaction_state: GcCompactionState, ) -> anyhow::Result<()> { self.gc_compaction_state - .store(Arc::new(Some(gc_compaction_state.clone()))); + .store(Some(Arc::new(gc_compaction_state.clone()))); self.remote_client .schedule_index_upload_for_gc_compaction_state_update(gc_compaction_state) } @@ -3429,7 +3434,10 @@ impl Timeline { } pub(crate) fn get_gc_compaction_state(&self) -> Option { - self.gc_compaction_state.load_full().as_ref().clone() + self.gc_compaction_state + .load() + .as_ref() + .map(|x| x.as_ref().clone()) } /// Creates and starts the wal receiver. diff --git a/test_runner/regress/test_feature_flag.py b/test_runner/regress/test_feature_flag.py index c6c192b6f1..6c1e3484fa 100644 --- a/test_runner/regress/test_feature_flag.py +++ b/test_runner/regress/test_feature_flag.py @@ -50,11 +50,15 @@ def test_feature_flag(neon_env_builder: NeonEnvBuilder): )["result"] ) + env.endpoints.create_start("main") # trigger basebackup env.pageserver.http_client().force_refresh_feature_flag(env.initial_tenant) # Check if the properties exist result = env.pageserver.http_client().evaluate_feature_flag_multivariate( env.initial_tenant, "test-feature-flag" ) + assert "tenant_remote_size_mb" in result["properties"] + assert "tenant_db_count_max" in result["properties"] + assert "tenant_rel_count_max" in result["properties"] assert "tenant_id" in result["properties"]