LCOV - code coverage report
Current view: top level - pageserver/src/tenant - size.rs (source / functions) Coverage Total Hit
Test: b4ae4c4857f9ef3e144e982a35ee23bc84c71983.info Lines: 40.3 % 588 237
Test Date: 2024-10-22 22:13:45 Functions: 23.9 % 71 17

            Line data    Source code
       1              : use std::cmp;
       2              : use std::collections::hash_map::Entry;
       3              : use std::collections::{HashMap, HashSet};
       4              : use std::sync::Arc;
       5              : 
       6              : use tenant_size_model::svg::SvgBranchKind;
       7              : use tokio::sync::oneshot::error::RecvError;
       8              : use tokio::sync::Semaphore;
       9              : use tokio_util::sync::CancellationToken;
      10              : 
      11              : use crate::context::RequestContext;
      12              : use crate::pgdatadir_mapping::CalculateLogicalSizeError;
      13              : 
      14              : use super::{GcError, LogicalSizeCalculationCause, Tenant};
      15              : use crate::tenant::{MaybeOffloaded, Timeline};
      16              : use utils::id::TimelineId;
      17              : use utils::lsn::Lsn;
      18              : 
      19              : use tracing::*;
      20              : 
      21              : use tenant_size_model::{Segment, StorageModel};
      22              : 
      23              : /// Inputs to the actual tenant sizing model
      24              : ///
      25              : /// Implements [`serde::Serialize`] but is not meant to be part of the public API, instead meant to
      26              : /// be a transferrable format between execution environments and developer.
      27              : ///
      28              : /// This tracks more information than the actual StorageModel that calculation
      29              : /// needs. We will convert this into a StorageModel when it's time to perform
      30              : /// the calculation.
      31              : ///
      32           12 : #[derive(Debug, serde::Serialize, serde::Deserialize)]
      33              : pub struct ModelInputs {
      34              :     pub segments: Vec<SegmentMeta>,
      35              :     pub timeline_inputs: Vec<TimelineInputs>,
      36              : }
      37              : 
      38              : /// A [`Segment`], with some extra information for display purposes
      39          112 : #[derive(Debug, serde::Serialize, serde::Deserialize)]
      40              : pub struct SegmentMeta {
      41              :     pub segment: Segment,
      42              :     pub timeline_id: TimelineId,
      43              :     pub kind: LsnKind,
      44              : }
      45              : 
      46            0 : #[derive(thiserror::Error, Debug)]
      47              : pub(crate) enum CalculateSyntheticSizeError {
      48              :     /// Something went wrong internally to the calculation of logical size at a particular branch point
      49              :     #[error("Failed to calculated logical size on timeline {timeline_id} at {lsn}: {error}")]
      50              :     LogicalSize {
      51              :         timeline_id: TimelineId,
      52              :         lsn: Lsn,
      53              :         error: CalculateLogicalSizeError,
      54              :     },
      55              : 
      56              :     /// Something went wrong internally when calculating GC parameters at start of size calculation
      57              :     #[error(transparent)]
      58              :     GcInfo(GcError),
      59              : 
      60              :     /// Totally unexpected errors, like panics joining a task
      61              :     #[error(transparent)]
      62              :     Fatal(anyhow::Error),
      63              : 
      64              :     /// Tenant shut down while calculating size
      65              :     #[error("Cancelled")]
      66              :     Cancelled,
      67              : }
      68              : 
      69              : impl From<GcError> for CalculateSyntheticSizeError {
      70            0 :     fn from(value: GcError) -> Self {
      71            0 :         match value {
      72              :             GcError::TenantCancelled | GcError::TimelineCancelled => {
      73            0 :                 CalculateSyntheticSizeError::Cancelled
      74              :             }
      75            0 :             other => CalculateSyntheticSizeError::GcInfo(other),
      76              :         }
      77            0 :     }
      78              : }
      79              : 
      80              : impl SegmentMeta {
      81            0 :     fn size_needed(&self) -> bool {
      82            0 :         match self.kind {
      83              :             LsnKind::BranchStart => {
      84              :                 // If we don't have a later GcCutoff point on this branch, and
      85              :                 // no ancestor, calculate size for the branch start point.
      86            0 :                 self.segment.needed && self.segment.parent.is_none()
      87              :             }
      88            0 :             LsnKind::BranchPoint => true,
      89            0 :             LsnKind::GcCutOff => true,
      90            0 :             LsnKind::BranchEnd => false,
      91            0 :             LsnKind::LeasePoint => true,
      92            0 :             LsnKind::LeaseStart => false,
      93            0 :             LsnKind::LeaseEnd => false,
      94              :         }
      95            0 :     }
      96              : }
      97              : 
      98              : #[derive(
      99           56 :     Debug, Clone, Copy, Eq, Ord, PartialEq, PartialOrd, serde::Serialize, serde::Deserialize,
     100              : )]
     101              : pub enum LsnKind {
     102              :     /// A timeline starting here
     103              :     BranchStart,
     104              :     /// A child timeline branches off from here
     105              :     BranchPoint,
     106              :     /// GC cutoff point
     107              :     GcCutOff,
     108              :     /// Last record LSN
     109              :     BranchEnd,
     110              :     /// A LSN lease is granted here.
     111              :     LeasePoint,
     112              :     /// A lease starts from here.
     113              :     LeaseStart,
     114              :     /// Last record LSN for the lease (should have the same LSN as the previous [`LsnKind::LeaseStart`]).
     115              :     LeaseEnd,
     116              : }
     117              : 
     118              : impl From<LsnKind> for SvgBranchKind {
     119            0 :     fn from(kind: LsnKind) -> Self {
     120            0 :         match kind {
     121            0 :             LsnKind::LeasePoint | LsnKind::LeaseStart | LsnKind::LeaseEnd => SvgBranchKind::Lease,
     122            0 :             _ => SvgBranchKind::Timeline,
     123              :         }
     124            0 :     }
     125              : }
     126              : 
     127              : /// Collect all relevant LSNs to the inputs. These will only be helpful in the serialized form as
     128              : /// part of [`ModelInputs`] from the HTTP api, explaining the inputs.
     129           64 : #[derive(Debug, serde::Serialize, serde::Deserialize)]
     130              : pub struct TimelineInputs {
     131              :     pub timeline_id: TimelineId,
     132              : 
     133              :     pub ancestor_id: Option<TimelineId>,
     134              : 
     135              :     ancestor_lsn: Lsn,
     136              :     last_record: Lsn,
     137              :     latest_gc_cutoff: Lsn,
     138              : 
     139              :     /// Cutoff point based on GC settings
     140              :     next_pitr_cutoff: Lsn,
     141              : 
     142              :     /// Cutoff point calculated from the user-supplied 'max_retention_period'
     143              :     retention_param_cutoff: Option<Lsn>,
     144              : 
     145              :     /// Lease points on the timeline
     146              :     lease_points: Vec<Lsn>,
     147              : }
     148              : 
     149              : /// Gathers the inputs for the tenant sizing model.
     150              : ///
     151              : /// Tenant size does not consider the latest state, but only the state until next_pitr_cutoff, which
     152              : /// is updated on-demand, during the start of this calculation and separate from the
     153              : /// [`TimelineInputs::latest_gc_cutoff`].
     154              : ///
     155              : /// For timelines in general:
     156              : ///
     157              : /// ```text
     158              : /// 0-----|---------|----|------------| · · · · · |·> lsn
     159              : ///   initdb_lsn  branchpoints*  next_pitr_cutoff  latest
     160              : /// ```
     161            0 : pub(super) async fn gather_inputs(
     162            0 :     tenant: &Tenant,
     163            0 :     limit: &Arc<Semaphore>,
     164            0 :     max_retention_period: Option<u64>,
     165            0 :     logical_size_cache: &mut HashMap<(TimelineId, Lsn), u64>,
     166            0 :     cause: LogicalSizeCalculationCause,
     167            0 :     cancel: &CancellationToken,
     168            0 :     ctx: &RequestContext,
     169            0 : ) -> Result<ModelInputs, CalculateSyntheticSizeError> {
     170            0 :     // refresh is needed to update [`timeline::GcCutoffs`]
     171            0 :     tenant.refresh_gc_info(cancel, ctx).await?;
     172              : 
     173              :     // Collect information about all the timelines
     174            0 :     let mut timelines = tenant.list_timelines();
     175            0 : 
     176            0 :     if timelines.is_empty() {
     177              :         // perhaps the tenant has just been created, and as such doesn't have any data yet
     178            0 :         return Ok(ModelInputs {
     179            0 :             segments: vec![],
     180            0 :             timeline_inputs: Vec::new(),
     181            0 :         });
     182            0 :     }
     183            0 : 
     184            0 :     // Filter out timelines that are not active
     185            0 :     //
     186            0 :     // There may be a race when a timeline is dropped,
     187            0 :     // but it is unlikely to cause any issues. In the worst case,
     188            0 :     // the calculation will error out.
     189            0 :     timelines.retain(|t| t.is_active());
     190            0 : 
     191            0 :     // Build a map of branch points.
     192            0 :     let mut branchpoints: HashMap<TimelineId, HashSet<Lsn>> = HashMap::new();
     193            0 :     for timeline in timelines.iter() {
     194            0 :         if let Some(ancestor_id) = timeline.get_ancestor_timeline_id() {
     195            0 :             branchpoints
     196            0 :                 .entry(ancestor_id)
     197            0 :                 .or_default()
     198            0 :                 .insert(timeline.get_ancestor_lsn());
     199            0 :         }
     200              :     }
     201              : 
     202              :     // These become the final result.
     203            0 :     let mut timeline_inputs = Vec::with_capacity(timelines.len());
     204            0 :     let mut segments: Vec<SegmentMeta> = Vec::new();
     205            0 : 
     206            0 :     //
     207            0 :     // Build Segments representing each timeline. As we do that, also remember
     208            0 :     // the branchpoints and branch startpoints in 'branchpoint_segments' and
     209            0 :     // 'branchstart_segments'
     210            0 :     //
     211            0 : 
     212            0 :     // BranchPoint segments of each timeline
     213            0 :     // (timeline, branchpoint LSN) -> segment_id
     214            0 :     let mut branchpoint_segments: HashMap<(TimelineId, Lsn), usize> = HashMap::new();
     215              : 
     216              :     // timeline, Branchpoint seg id, (ancestor, ancestor LSN)
     217              :     type BranchStartSegment = (TimelineId, usize, Option<(TimelineId, Lsn)>);
     218            0 :     let mut branchstart_segments: Vec<BranchStartSegment> = Vec::new();
     219              : 
     220            0 :     for timeline in timelines.iter() {
     221            0 :         let timeline_id = timeline.timeline_id;
     222            0 :         let last_record_lsn = timeline.get_last_record_lsn();
     223            0 :         let ancestor_lsn = timeline.get_ancestor_lsn();
     224            0 : 
     225            0 :         // there's a race between the update (holding tenant.gc_lock) and this read but it
     226            0 :         // might not be an issue, because it's not for Timeline::gc
     227            0 :         let gc_info = timeline.gc_info.read().unwrap();
     228            0 : 
     229            0 :         // similar to gc, but Timeline::get_latest_gc_cutoff_lsn() will not be updated before a
     230            0 :         // new gc run, which we have no control over. however differently from `Timeline::gc`
     231            0 :         // we don't consider the `Timeline::disk_consistent_lsn` at all, because we are not
     232            0 :         // actually removing files.
     233            0 :         //
     234            0 :         // We only consider [`timeline::GcCutoffs::time`], and not [`timeline::GcCutoffs::space`], because from
     235            0 :         // a user's perspective they have only requested retention up to the time bound (pitr_cutoff), rather
     236            0 :         // than our internal space cutoff.  This means that if someone drops a database and waits for their
     237            0 :         // PITR interval, they will see synthetic size decrease, even if we are still storing data inside
     238            0 :         // the space cutoff.
     239            0 :         let mut next_pitr_cutoff = gc_info.cutoffs.time;
     240              : 
     241              :         // If the caller provided a shorter retention period, use that instead of the GC cutoff.
     242            0 :         let retention_param_cutoff = if let Some(max_retention_period) = max_retention_period {
     243            0 :             let param_cutoff = Lsn(last_record_lsn.0.saturating_sub(max_retention_period));
     244            0 :             if next_pitr_cutoff < param_cutoff {
     245            0 :                 next_pitr_cutoff = param_cutoff;
     246            0 :             }
     247            0 :             Some(param_cutoff)
     248              :         } else {
     249            0 :             None
     250              :         };
     251              : 
     252            0 :         let lease_points = gc_info
     253            0 :             .leases
     254            0 :             .keys()
     255            0 :             .filter(|&&lsn| lsn > ancestor_lsn)
     256            0 :             .copied()
     257            0 :             .collect::<Vec<_>>();
     258            0 : 
     259            0 :         // next_pitr_cutoff in parent branch are not of interest (right now at least), nor do we
     260            0 :         // want to query any logical size before initdb_lsn.
     261            0 :         let branch_start_lsn = cmp::max(ancestor_lsn, timeline.initdb_lsn);
     262            0 : 
     263            0 :         // Build "interesting LSNs" on this timeline
     264            0 :         let mut lsns: Vec<(Lsn, LsnKind)> = gc_info
     265            0 :             .retain_lsns
     266            0 :             .iter()
     267            0 :             .filter(|(lsn, _child_id, is_offloaded)| {
     268            0 :                 lsn > &ancestor_lsn && *is_offloaded == MaybeOffloaded::No
     269            0 :             })
     270            0 :             .copied()
     271            0 :             // this assumes there are no other retain_lsns than the branchpoints
     272            0 :             .map(|(lsn, _child_id, _is_offloaded)| (lsn, LsnKind::BranchPoint))
     273            0 :             .collect::<Vec<_>>();
     274            0 : 
     275            0 :         lsns.extend(lease_points.iter().map(|&lsn| (lsn, LsnKind::LeasePoint)));
     276            0 : 
     277            0 :         drop(gc_info);
     278              : 
     279              :         // Add branch points we collected earlier, just in case there were any that were
     280              :         // not present in retain_lsns. We will remove any duplicates below later.
     281            0 :         if let Some(this_branchpoints) = branchpoints.get(&timeline_id) {
     282            0 :             lsns.extend(
     283            0 :                 this_branchpoints
     284            0 :                     .iter()
     285            0 :                     .map(|lsn| (*lsn, LsnKind::BranchPoint)),
     286            0 :             )
     287            0 :         }
     288              : 
     289              :         // Add a point for the PITR cutoff
     290            0 :         let branch_start_needed = next_pitr_cutoff <= branch_start_lsn;
     291            0 :         if !branch_start_needed {
     292            0 :             lsns.push((next_pitr_cutoff, LsnKind::GcCutOff));
     293            0 :         }
     294              : 
     295            0 :         lsns.sort_unstable();
     296            0 :         lsns.dedup();
     297            0 : 
     298            0 :         //
     299            0 :         // Create Segments for the interesting points.
     300            0 :         //
     301            0 : 
     302            0 :         // Timeline start point
     303            0 :         let ancestor = timeline
     304            0 :             .get_ancestor_timeline_id()
     305            0 :             .map(|ancestor_id| (ancestor_id, ancestor_lsn));
     306            0 :         branchstart_segments.push((timeline_id, segments.len(), ancestor));
     307            0 :         segments.push(SegmentMeta {
     308            0 :             segment: Segment {
     309            0 :                 parent: None, // filled in later
     310            0 :                 lsn: branch_start_lsn.0,
     311            0 :                 size: None, // filled in later
     312            0 :                 needed: branch_start_needed,
     313            0 :             },
     314            0 :             timeline_id: timeline.timeline_id,
     315            0 :             kind: LsnKind::BranchStart,
     316            0 :         });
     317            0 : 
     318            0 :         // GC cutoff point, and any branch points, i.e. points where
     319            0 :         // other timelines branch off from this timeline.
     320            0 :         let mut parent = segments.len() - 1;
     321            0 :         for (lsn, kind) in lsns {
     322            0 :             if kind == LsnKind::BranchPoint {
     323            0 :                 branchpoint_segments.insert((timeline_id, lsn), segments.len());
     324            0 :             }
     325              : 
     326            0 :             segments.push(SegmentMeta {
     327            0 :                 segment: Segment {
     328            0 :                     parent: Some(parent),
     329            0 :                     lsn: lsn.0,
     330            0 :                     size: None,
     331            0 :                     needed: lsn > next_pitr_cutoff,
     332            0 :                 },
     333            0 :                 timeline_id: timeline.timeline_id,
     334            0 :                 kind,
     335            0 :             });
     336            0 : 
     337            0 :             parent = segments.len() - 1;
     338            0 : 
     339            0 :             if kind == LsnKind::LeasePoint {
     340            0 :                 // Needs `LeaseStart` and `LeaseEnd` as well to model lease as a read-only branch that never writes data
     341            0 :                 // (i.e. it's lsn has not advanced from ancestor_lsn), and therefore the three segments have the same LSN
     342            0 :                 // value. Without the other two segments, the calculation code would not count the leased LSN as a point
     343            0 :                 // to be retained.
     344            0 :                 // Did not use `BranchStart` or `BranchEnd` so we can differentiate branches and leases during debug.
     345            0 :                 //
     346            0 :                 // Alt Design: rewrite the entire calculation code to be independent of timeline id. Both leases and
     347            0 :                 // branch points can be given a synthetic id so we can unite them.
     348            0 :                 let mut lease_parent = parent;
     349            0 : 
     350            0 :                 // Start of a lease.
     351            0 :                 segments.push(SegmentMeta {
     352            0 :                     segment: Segment {
     353            0 :                         parent: Some(lease_parent),
     354            0 :                         lsn: lsn.0,
     355            0 :                         size: None,                     // Filled in later, if necessary
     356            0 :                         needed: lsn > next_pitr_cutoff, // only needed if the point is within rentention.
     357            0 :                     },
     358            0 :                     timeline_id: timeline.timeline_id,
     359            0 :                     kind: LsnKind::LeaseStart,
     360            0 :                 });
     361            0 :                 lease_parent += 1;
     362            0 : 
     363            0 :                 // End of the lease.
     364            0 :                 segments.push(SegmentMeta {
     365            0 :                     segment: Segment {
     366            0 :                         parent: Some(lease_parent),
     367            0 :                         lsn: lsn.0,
     368            0 :                         size: None,   // Filled in later, if necessary
     369            0 :                         needed: true, // everything at the lease LSN must be readable => is needed
     370            0 :                     },
     371            0 :                     timeline_id: timeline.timeline_id,
     372            0 :                     kind: LsnKind::LeaseEnd,
     373            0 :                 });
     374            0 :             }
     375              :         }
     376              : 
     377              :         // Current end of the timeline
     378            0 :         segments.push(SegmentMeta {
     379            0 :             segment: Segment {
     380            0 :                 parent: Some(parent),
     381            0 :                 lsn: last_record_lsn.0,
     382            0 :                 size: None, // Filled in later, if necessary
     383            0 :                 needed: true,
     384            0 :             },
     385            0 :             timeline_id: timeline.timeline_id,
     386            0 :             kind: LsnKind::BranchEnd,
     387            0 :         });
     388            0 : 
     389            0 :         timeline_inputs.push(TimelineInputs {
     390            0 :             timeline_id: timeline.timeline_id,
     391            0 :             ancestor_id: timeline.get_ancestor_timeline_id(),
     392            0 :             ancestor_lsn,
     393            0 :             last_record: last_record_lsn,
     394            0 :             // this is not used above, because it might not have updated recently enough
     395            0 :             latest_gc_cutoff: *timeline.get_latest_gc_cutoff_lsn(),
     396            0 :             next_pitr_cutoff,
     397            0 :             retention_param_cutoff,
     398            0 :             lease_points,
     399            0 :         });
     400              :     }
     401              : 
     402              :     // We now have all segments from the timelines in 'segments'. The timelines
     403              :     // haven't been linked to each other yet, though. Do that.
     404            0 :     for (_timeline_id, seg_id, ancestor) in branchstart_segments {
     405              :         // Look up the branch point
     406            0 :         if let Some(ancestor) = ancestor {
     407            0 :             let parent_id = *branchpoint_segments.get(&ancestor).unwrap();
     408            0 :             segments[seg_id].segment.parent = Some(parent_id);
     409            0 :         }
     410              :     }
     411              : 
     412              :     // We left the 'size' field empty in all of the Segments so far.
     413              :     // Now find logical sizes for all of the points that might need or benefit from them.
     414            0 :     fill_logical_sizes(
     415            0 :         &timelines,
     416            0 :         &mut segments,
     417            0 :         limit,
     418            0 :         logical_size_cache,
     419            0 :         cause,
     420            0 :         ctx,
     421            0 :     )
     422            0 :     .await?;
     423              : 
     424            0 :     if tenant.cancel.is_cancelled() {
     425              :         // If we're shutting down, return an error rather than a sparse result that might include some
     426              :         // timelines from before we started shutting down
     427            0 :         return Err(CalculateSyntheticSizeError::Cancelled);
     428            0 :     }
     429            0 : 
     430            0 :     Ok(ModelInputs {
     431            0 :         segments,
     432            0 :         timeline_inputs,
     433            0 :     })
     434            0 : }
     435              : 
     436              : /// Augment 'segments' with logical sizes
     437              : ///
     438              : /// This will leave segments' sizes as None if the Timeline associated with the segment is deleted concurrently
     439              : /// (i.e. we cannot read its logical size at a particular LSN).
     440            0 : async fn fill_logical_sizes(
     441            0 :     timelines: &[Arc<Timeline>],
     442            0 :     segments: &mut [SegmentMeta],
     443            0 :     limit: &Arc<Semaphore>,
     444            0 :     logical_size_cache: &mut HashMap<(TimelineId, Lsn), u64>,
     445            0 :     cause: LogicalSizeCalculationCause,
     446            0 :     ctx: &RequestContext,
     447            0 : ) -> Result<(), CalculateSyntheticSizeError> {
     448            0 :     let timeline_hash: HashMap<TimelineId, Arc<Timeline>> = HashMap::from_iter(
     449            0 :         timelines
     450            0 :             .iter()
     451            0 :             .map(|timeline| (timeline.timeline_id, Arc::clone(timeline))),
     452            0 :     );
     453            0 : 
     454            0 :     // record the used/inserted cache keys here, to remove extras not to start leaking
     455            0 :     // after initial run the cache should be quite stable, but live timelines will eventually
     456            0 :     // require new lsns to be inspected.
     457            0 :     let mut sizes_needed = HashMap::<(TimelineId, Lsn), Option<u64>>::new();
     458            0 : 
     459            0 :     // with joinset, on drop, all of the tasks will just be de-scheduled, which we can use to
     460            0 :     // our advantage with `?` error handling.
     461            0 :     let mut joinset = tokio::task::JoinSet::new();
     462              : 
     463              :     // For each point that would benefit from having a logical size available,
     464              :     // spawn a Task to fetch it, unless we have it cached already.
     465            0 :     for seg in segments.iter() {
     466            0 :         if !seg.size_needed() {
     467            0 :             continue;
     468            0 :         }
     469            0 : 
     470            0 :         let timeline_id = seg.timeline_id;
     471            0 :         let lsn = Lsn(seg.segment.lsn);
     472              : 
     473            0 :         if let Entry::Vacant(e) = sizes_needed.entry((timeline_id, lsn)) {
     474            0 :             let cached_size = logical_size_cache.get(&(timeline_id, lsn)).cloned();
     475            0 :             if cached_size.is_none() {
     476            0 :                 let timeline = Arc::clone(timeline_hash.get(&timeline_id).unwrap());
     477            0 :                 let parallel_size_calcs = Arc::clone(limit);
     478            0 :                 let ctx = ctx.attached_child();
     479            0 :                 joinset.spawn(
     480            0 :                     calculate_logical_size(parallel_size_calcs, timeline, lsn, cause, ctx)
     481            0 :                         .in_current_span(),
     482            0 :                 );
     483            0 :             }
     484            0 :             e.insert(cached_size);
     485            0 :         }
     486              :     }
     487              : 
     488              :     // Perform the size lookups
     489            0 :     let mut have_any_error = None;
     490            0 :     while let Some(res) = joinset.join_next().await {
     491              :         // each of these come with Result<anyhow::Result<_>, JoinError>
     492              :         // because of spawn + spawn_blocking
     493            0 :         match res {
     494            0 :             Err(join_error) if join_error.is_cancelled() => {
     495            0 :                 unreachable!("we are not cancelling any of the futures, nor should be");
     496              :             }
     497            0 :             Err(join_error) => {
     498            0 :                 // cannot really do anything, as this panic is likely a bug
     499            0 :                 error!("task that calls spawn_ondemand_logical_size_calculation panicked: {join_error:#}");
     500              : 
     501            0 :                 have_any_error = Some(CalculateSyntheticSizeError::Fatal(
     502            0 :                     anyhow::anyhow!(join_error)
     503            0 :                         .context("task that calls spawn_ondemand_logical_size_calculation"),
     504            0 :                 ));
     505              :             }
     506            0 :             Ok(Err(recv_result_error)) => {
     507            0 :                 // cannot really do anything, as this panic is likely a bug
     508            0 :                 error!("failed to receive logical size query result: {recv_result_error:#}");
     509            0 :                 have_any_error = Some(CalculateSyntheticSizeError::Fatal(
     510            0 :                     anyhow::anyhow!(recv_result_error)
     511            0 :                         .context("Receiving logical size query result"),
     512            0 :                 ));
     513              :             }
     514            0 :             Ok(Ok(TimelineAtLsnSizeResult(timeline, lsn, Err(error)))) => {
     515            0 :                 if matches!(error, CalculateLogicalSizeError::Cancelled) {
     516              :                     // Skip this: it's okay if one timeline among many is shutting down while we
     517              :                     // calculate inputs for the overall tenant.
     518            0 :                     continue;
     519              :                 } else {
     520            0 :                     warn!(
     521            0 :                         timeline_id=%timeline.timeline_id,
     522            0 :                         "failed to calculate logical size at {lsn}: {error:#}"
     523              :                     );
     524            0 :                     have_any_error = Some(CalculateSyntheticSizeError::LogicalSize {
     525            0 :                         timeline_id: timeline.timeline_id,
     526            0 :                         lsn,
     527            0 :                         error,
     528            0 :                     });
     529              :                 }
     530              :             }
     531            0 :             Ok(Ok(TimelineAtLsnSizeResult(timeline, lsn, Ok(size)))) => {
     532            0 :                 debug!(timeline_id=%timeline.timeline_id, %lsn, size, "size calculated");
     533              : 
     534            0 :                 logical_size_cache.insert((timeline.timeline_id, lsn), size);
     535            0 :                 sizes_needed.insert((timeline.timeline_id, lsn), Some(size));
     536              :             }
     537              :         }
     538              :     }
     539              : 
     540              :     // prune any keys not needed anymore; we record every used key and added key.
     541            0 :     logical_size_cache.retain(|key, _| sizes_needed.contains_key(key));
     542              : 
     543            0 :     if let Some(error) = have_any_error {
     544              :         // we cannot complete this round, because we are missing data.
     545              :         // we have however cached all we were able to request calculation on.
     546            0 :         return Err(error);
     547            0 :     }
     548              : 
     549              :     // Insert the looked up sizes to the Segments
     550            0 :     for seg in segments.iter_mut() {
     551            0 :         if !seg.size_needed() {
     552            0 :             continue;
     553            0 :         }
     554            0 : 
     555            0 :         let timeline_id = seg.timeline_id;
     556            0 :         let lsn = Lsn(seg.segment.lsn);
     557              : 
     558            0 :         if let Some(Some(size)) = sizes_needed.get(&(timeline_id, lsn)) {
     559            0 :             seg.segment.size = Some(*size);
     560            0 :         }
     561              :     }
     562            0 :     Ok(())
     563            0 : }
     564              : 
     565              : impl ModelInputs {
     566            4 :     pub fn calculate_model(&self) -> tenant_size_model::StorageModel {
     567            4 :         // Convert SegmentMetas into plain Segments
     568            4 :         StorageModel {
     569            4 :             segments: self
     570            4 :                 .segments
     571            4 :                 .iter()
     572           28 :                 .map(|seg| seg.segment.clone())
     573            4 :                 .collect(),
     574            4 :         }
     575            4 :     }
     576              : 
     577              :     // calculate total project size
     578            2 :     pub fn calculate(&self) -> u64 {
     579            2 :         let storage = self.calculate_model();
     580            2 :         let sizes = storage.calculate();
     581            2 :         sizes.total_size
     582            2 :     }
     583              : }
     584              : 
     585              : /// Newtype around the tuple that carries the timeline at lsn logical size calculation.
     586              : struct TimelineAtLsnSizeResult(
     587              :     Arc<crate::tenant::Timeline>,
     588              :     utils::lsn::Lsn,
     589              :     Result<u64, CalculateLogicalSizeError>,
     590              : );
     591              : 
     592            0 : #[instrument(skip_all, fields(timeline_id=%timeline.timeline_id, lsn=%lsn))]
     593              : async fn calculate_logical_size(
     594              :     limit: Arc<tokio::sync::Semaphore>,
     595              :     timeline: Arc<crate::tenant::Timeline>,
     596              :     lsn: utils::lsn::Lsn,
     597              :     cause: LogicalSizeCalculationCause,
     598              :     ctx: RequestContext,
     599              : ) -> Result<TimelineAtLsnSizeResult, RecvError> {
     600              :     let _permit = tokio::sync::Semaphore::acquire_owned(limit)
     601              :         .await
     602              :         .expect("global semaphore should not had been closed");
     603              : 
     604              :     let size_res = timeline
     605              :         .spawn_ondemand_logical_size_calculation(lsn, cause, ctx)
     606              :         .instrument(info_span!("spawn_ondemand_logical_size_calculation"))
     607              :         .await?;
     608              :     Ok(TimelineAtLsnSizeResult(timeline, lsn, size_res))
     609              : }
     610              : 
     611              : #[test]
     612            2 : fn verify_size_for_multiple_branches() {
     613            2 :     // this is generated from integration test test_tenant_size_with_multiple_branches, but this way
     614            2 :     // it has the stable lsn's
     615            2 :     //
     616            2 :     // The timeline_inputs don't participate in the size calculation, and are here just to explain
     617            2 :     // the inputs.
     618            2 :     let doc = r#"
     619            2 : {
     620            2 :   "segments": [
     621            2 :     {
     622            2 :       "segment": {
     623            2 :         "parent": 9,
     624            2 :         "lsn": 26033560,
     625            2 :         "size": null,
     626            2 :         "needed": false
     627            2 :       },
     628            2 :       "timeline_id": "20b129c9b50cff7213e6503a31b2a5ce",
     629            2 :       "kind": "BranchStart"
     630            2 :     },
     631            2 :     {
     632            2 :       "segment": {
     633            2 :         "parent": 0,
     634            2 :         "lsn": 35720400,
     635            2 :         "size": 25206784,
     636            2 :         "needed": false
     637            2 :       },
     638            2 :       "timeline_id": "20b129c9b50cff7213e6503a31b2a5ce",
     639            2 :       "kind": "GcCutOff"
     640            2 :     },
     641            2 :     {
     642            2 :       "segment": {
     643            2 :         "parent": 1,
     644            2 :         "lsn": 35851472,
     645            2 :         "size": null,
     646            2 :         "needed": true
     647            2 :       },
     648            2 :       "timeline_id": "20b129c9b50cff7213e6503a31b2a5ce",
     649            2 :       "kind": "BranchEnd"
     650            2 :     },
     651            2 :     {
     652            2 :       "segment": {
     653            2 :         "parent": 7,
     654            2 :         "lsn": 24566168,
     655            2 :         "size": null,
     656            2 :         "needed": false
     657            2 :       },
     658            2 :       "timeline_id": "454626700469f0a9914949b9d018e876",
     659            2 :       "kind": "BranchStart"
     660            2 :     },
     661            2 :     {
     662            2 :       "segment": {
     663            2 :         "parent": 3,
     664            2 :         "lsn": 25261936,
     665            2 :         "size": 26050560,
     666            2 :         "needed": false
     667            2 :       },
     668            2 :       "timeline_id": "454626700469f0a9914949b9d018e876",
     669            2 :       "kind": "GcCutOff"
     670            2 :     },
     671            2 :     {
     672            2 :       "segment": {
     673            2 :         "parent": 4,
     674            2 :         "lsn": 25393008,
     675            2 :         "size": null,
     676            2 :         "needed": true
     677            2 :       },
     678            2 :       "timeline_id": "454626700469f0a9914949b9d018e876",
     679            2 :       "kind": "BranchEnd"
     680            2 :     },
     681            2 :     {
     682            2 :       "segment": {
     683            2 :         "parent": null,
     684            2 :         "lsn": 23694408,
     685            2 :         "size": null,
     686            2 :         "needed": false
     687            2 :       },
     688            2 :       "timeline_id": "cb5e3cbe60a4afc00d01880e1a37047f",
     689            2 :       "kind": "BranchStart"
     690            2 :     },
     691            2 :     {
     692            2 :       "segment": {
     693            2 :         "parent": 6,
     694            2 :         "lsn": 24566168,
     695            2 :         "size": 25739264,
     696            2 :         "needed": false
     697            2 :       },
     698            2 :       "timeline_id": "cb5e3cbe60a4afc00d01880e1a37047f",
     699            2 :       "kind": "BranchPoint"
     700            2 :     },
     701            2 :     {
     702            2 :       "segment": {
     703            2 :         "parent": 7,
     704            2 :         "lsn": 25902488,
     705            2 :         "size": 26402816,
     706            2 :         "needed": false
     707            2 :       },
     708            2 :       "timeline_id": "cb5e3cbe60a4afc00d01880e1a37047f",
     709            2 :       "kind": "GcCutOff"
     710            2 :     },
     711            2 :     {
     712            2 :       "segment": {
     713            2 :         "parent": 8,
     714            2 :         "lsn": 26033560,
     715            2 :         "size": 26468352,
     716            2 :         "needed": true
     717            2 :       },
     718            2 :       "timeline_id": "cb5e3cbe60a4afc00d01880e1a37047f",
     719            2 :       "kind": "BranchPoint"
     720            2 :     },
     721            2 :     {
     722            2 :       "segment": {
     723            2 :         "parent": 9,
     724            2 :         "lsn": 26033560,
     725            2 :         "size": null,
     726            2 :         "needed": true
     727            2 :       },
     728            2 :       "timeline_id": "cb5e3cbe60a4afc00d01880e1a37047f",
     729            2 :       "kind": "BranchEnd"
     730            2 :     }
     731            2 :   ],
     732            2 :   "timeline_inputs": [
     733            2 :     {
     734            2 :       "timeline_id": "20b129c9b50cff7213e6503a31b2a5ce",
     735            2 :       "ancestor_lsn": "0/18D3D98",
     736            2 :       "last_record": "0/2230CD0",
     737            2 :       "latest_gc_cutoff": "0/1698C48",
     738            2 :       "next_pitr_cutoff": "0/2210CD0",
     739            2 :       "retention_param_cutoff": null,
     740            2 :       "lease_points": []
     741            2 :     },
     742            2 :     {
     743            2 :       "timeline_id": "454626700469f0a9914949b9d018e876",
     744            2 :       "ancestor_lsn": "0/176D998",
     745            2 :       "last_record": "0/1837770",
     746            2 :       "latest_gc_cutoff": "0/1698C48",
     747            2 :       "next_pitr_cutoff": "0/1817770",
     748            2 :       "retention_param_cutoff": null,
     749            2 :       "lease_points": []
     750            2 :     },
     751            2 :     {
     752            2 :       "timeline_id": "cb5e3cbe60a4afc00d01880e1a37047f",
     753            2 :       "ancestor_lsn": "0/0",
     754            2 :       "last_record": "0/18D3D98",
     755            2 :       "latest_gc_cutoff": "0/1698C48",
     756            2 :       "next_pitr_cutoff": "0/18B3D98",
     757            2 :       "retention_param_cutoff": null,
     758            2 :       "lease_points": []
     759            2 :     }
     760            2 :   ]
     761            2 : }
     762            2 : "#;
     763            2 :     let inputs: ModelInputs = serde_json::from_str(doc).unwrap();
     764            2 : 
     765            2 :     assert_eq!(inputs.calculate(), 37_851_408);
     766            2 : }
     767              : 
     768              : #[test]
     769            2 : fn verify_size_for_one_branch() {
     770            2 :     let doc = r#"
     771            2 : {
     772            2 :   "segments": [
     773            2 :     {
     774            2 :       "segment": {
     775            2 :         "parent": null,
     776            2 :         "lsn": 0,
     777            2 :         "size": null,
     778            2 :         "needed": false
     779            2 :       },
     780            2 :       "timeline_id": "f15ae0cf21cce2ba27e4d80c6709a6cd",
     781            2 :       "kind": "BranchStart"
     782            2 :     },
     783            2 :     {
     784            2 :       "segment": {
     785            2 :         "parent": 0,
     786            2 :         "lsn": 305547335776,
     787            2 :         "size": 220054675456,
     788            2 :         "needed": false
     789            2 :       },
     790            2 :       "timeline_id": "f15ae0cf21cce2ba27e4d80c6709a6cd",
     791            2 :       "kind": "GcCutOff"
     792            2 :     },
     793            2 :     {
     794            2 :       "segment": {
     795            2 :         "parent": 1,
     796            2 :         "lsn": 305614444640,
     797            2 :         "size": null,
     798            2 :         "needed": true
     799            2 :       },
     800            2 :       "timeline_id": "f15ae0cf21cce2ba27e4d80c6709a6cd",
     801            2 :       "kind": "BranchEnd"
     802            2 :     }
     803            2 :   ],
     804            2 :   "timeline_inputs": [
     805            2 :     {
     806            2 :       "timeline_id": "f15ae0cf21cce2ba27e4d80c6709a6cd",
     807            2 :       "ancestor_lsn": "0/0",
     808            2 :       "last_record": "47/280A5860",
     809            2 :       "latest_gc_cutoff": "47/240A5860",
     810            2 :       "next_pitr_cutoff": "47/240A5860",
     811            2 :       "retention_param_cutoff": "0/0",
     812            2 :       "lease_points": []
     813            2 :     }
     814            2 :   ]
     815            2 : }"#;
     816            2 : 
     817            2 :     let model: ModelInputs = serde_json::from_str(doc).unwrap();
     818            2 : 
     819            2 :     let res = model.calculate_model().calculate();
     820            2 : 
     821            2 :     println!("calculated synthetic size: {}", res.total_size);
     822            2 :     println!("result: {:?}", serde_json::to_string(&res.segments));
     823              : 
     824              :     use utils::lsn::Lsn;
     825            2 :     let latest_gc_cutoff_lsn: Lsn = "47/240A5860".parse().unwrap();
     826            2 :     let last_lsn: Lsn = "47/280A5860".parse().unwrap();
     827            2 :     println!(
     828            2 :         "latest_gc_cutoff lsn 47/240A5860 is {}, last_lsn lsn 47/280A5860 is {}",
     829            2 :         u64::from(latest_gc_cutoff_lsn),
     830            2 :         u64::from(last_lsn)
     831            2 :     );
     832            2 :     assert_eq!(res.total_size, 220121784320);
     833            2 : }
        

Generated by: LCOV version 2.1-beta