LCOV - code coverage report
Current view: top level - libs/tenant_size_model/src - calculation.rs (source / functions) Coverage Total Hit
Test: 1e20c4f2b28aa592527961bb32170ebbd2c9172f.info Lines: 98.2 % 111 109
Test Date: 2025-07-16 12:29:03 Functions: 100.0 % 3 3

            Line data    Source code
       1              : use crate::{SegmentMethod, SegmentSizeResult, SizeResult, StorageModel};
       2              : 
       3              : //
       4              : //                 *-g--*---D--->
       5              : //                /
       6              : //               /
       7              : //              /                 *---b----*-B--->
       8              : //             /                 /
       9              : //            /                 /
      10              : //      -----*--e---*-----f----* C
      11              : //           E                  \
      12              : //                               \
      13              : //                                *--a---*---A-->
      14              : //
      15              : // If A and B need to be retained, is it cheaper to store
      16              : // snapshot at C+a+b, or snapshots at A and B ?
      17              : //
      18              : // If D also needs to be retained, which is cheaper:
      19              : //
      20              : // 1. E+g+e+f+a+b
      21              : // 2. D+C+a+b
      22              : // 3. D+A+B
      23              : 
      24              : /// `Segment` which has had its size calculated.
      25              : #[derive(Clone, Debug)]
      26              : struct SegmentSize {
      27              :     method: SegmentMethod,
      28              : 
      29              :     // calculated size of this subtree, using this method
      30              :     accum_size: u64,
      31              : 
      32              :     seg_id: usize,
      33              :     children: Vec<SegmentSize>,
      34              : }
      35              : 
      36              : struct SizeAlternatives {
      37              :     /// cheapest alternative if parent is available.
      38              :     incremental: SegmentSize,
      39              : 
      40              :     /// cheapest alternative if parent node is not available
      41              :     non_incremental: Option<SegmentSize>,
      42              : }
      43              : 
      44              : impl StorageModel {
      45            8 :     pub fn calculate(&self) -> SizeResult {
      46              :         // Build adjacency list. 'child_list' is indexed by segment id. Each entry
      47              :         // contains a list of all child segments of the segment.
      48            8 :         let mut roots: Vec<usize> = Vec::new();
      49            8 :         let mut child_list: Vec<Vec<usize>> = Vec::new();
      50            8 :         child_list.resize(self.segments.len(), Vec::new());
      51              : 
      52           71 :         for (seg_id, seg) in self.segments.iter().enumerate() {
      53           71 :             if let Some(parent_id) = seg.parent {
      54           63 :                 child_list[parent_id].push(seg_id);
      55           63 :             } else {
      56            8 :                 roots.push(seg_id);
      57            8 :             }
      58              :         }
      59              : 
      60            8 :         let mut segment_results = Vec::new();
      61            8 :         segment_results.resize(
      62            8 :             self.segments.len(),
      63            8 :             SegmentSizeResult {
      64            8 :                 method: SegmentMethod::Skipped,
      65            8 :                 accum_size: 0,
      66            8 :             },
      67              :         );
      68              : 
      69            8 :         let mut total_size = 0;
      70           16 :         for root in roots {
      71            8 :             if let Some(selected) = self.size_here(root, &child_list).non_incremental {
      72            8 :                 StorageModel::fill_selected_sizes(&selected, &mut segment_results);
      73            8 :                 total_size += selected.accum_size;
      74            8 :             } else {
      75            0 :                 // Couldn't find any way to get this root. Error?
      76            0 :             }
      77              :         }
      78              : 
      79            8 :         SizeResult {
      80            8 :             // If total_size is 0, it means that the tenant has all timelines offloaded; we need to report 1
      81            8 :             // here so that the data point shows up in the s3 files.
      82            8 :             total_size: total_size.max(1),
      83            8 :             segments: segment_results,
      84            8 :         }
      85            8 :     }
      86              : 
      87           71 :     fn fill_selected_sizes(selected: &SegmentSize, result: &mut Vec<SegmentSizeResult>) {
      88           71 :         result[selected.seg_id] = SegmentSizeResult {
      89           71 :             method: selected.method,
      90           71 :             accum_size: selected.accum_size,
      91           71 :         };
      92              :         // recurse to children
      93           71 :         for child in selected.children.iter() {
      94           63 :             StorageModel::fill_selected_sizes(child, result);
      95           63 :         }
      96           71 :     }
      97              : 
      98              :     //
      99              :     // This is the core of the sizing calculation.
     100              :     //
     101              :     // This is a recursive function, that for each Segment calculates the best way
     102              :     // to reach all the Segments that are marked as needed in this subtree, under two
     103              :     // different conditions:
     104              :     // a) when the parent of this segment is available (as a snaphot or through WAL), and
     105              :     // b) when the parent of this segment is not available.
     106              :     //
     107           71 :     fn size_here(&self, seg_id: usize, child_list: &Vec<Vec<usize>>) -> SizeAlternatives {
     108           71 :         let seg = &self.segments[seg_id];
     109              :         // First figure out the best way to get each child
     110           71 :         let mut children = Vec::new();
     111           71 :         for child_id in &child_list[seg_id] {
     112           63 :             children.push(self.size_here(*child_id, child_list))
     113              :         }
     114              : 
     115              :         // Method 1. If this node is not needed, we can skip it as long as we
     116              :         // take snapshots later in each sub-tree
     117           71 :         let snapshot_later = if !seg.needed {
     118           53 :             let mut snapshot_later = SegmentSize {
     119           53 :                 seg_id,
     120           53 :                 method: SegmentMethod::Skipped,
     121           53 :                 accum_size: 0,
     122           53 :                 children: Vec::new(),
     123           53 :             };
     124              : 
     125           53 :             let mut possible = true;
     126           54 :             for child in children.iter() {
     127           54 :                 if let Some(non_incremental) = &child.non_incremental {
     128           40 :                     snapshot_later.accum_size += non_incremental.accum_size;
     129           40 :                     snapshot_later.children.push(non_incremental.clone())
     130              :                 } else {
     131           14 :                     possible = false;
     132           14 :                     break;
     133              :                 }
     134              :             }
     135           53 :             if possible { Some(snapshot_later) } else { None }
     136              :         } else {
     137           18 :             None
     138              :         };
     139              : 
     140              :         // Method 2. Get a snapshot here. This assumed to be possible, if the 'size' of
     141              :         // this Segment was given.
     142           71 :         let snapshot_here = if !seg.needed || seg.parent.is_none() {
     143           53 :             if let Some(snapshot_size) = seg.size {
     144           49 :                 let mut snapshot_here = SegmentSize {
     145           49 :                     seg_id,
     146           49 :                     method: SegmentMethod::SnapshotHere,
     147           49 :                     accum_size: snapshot_size,
     148           49 :                     children: Vec::new(),
     149           49 :                 };
     150           57 :                 for child in children.iter() {
     151           57 :                     snapshot_here.accum_size += child.incremental.accum_size;
     152           57 :                     snapshot_here.children.push(child.incremental.clone())
     153              :                 }
     154           49 :                 Some(snapshot_here)
     155              :             } else {
     156            4 :                 None
     157              :             }
     158              :         } else {
     159           18 :             None
     160              :         };
     161              : 
     162              :         // Method 3. Use WAL to get here from parent
     163           71 :         let wal_here = {
     164           71 :             let mut wal_here = SegmentSize {
     165           71 :                 seg_id,
     166           71 :                 method: SegmentMethod::Wal,
     167           71 :                 accum_size: if let Some(parent_id) = seg.parent {
     168           63 :                     seg.lsn - self.segments[parent_id].lsn
     169              :                 } else {
     170            8 :                     0
     171              :                 },
     172           71 :                 children: Vec::new(),
     173              :             };
     174          134 :             for child in children {
     175           63 :                 wal_here.accum_size += child.incremental.accum_size;
     176           63 :                 wal_here.children.push(child.incremental)
     177              :             }
     178           71 :             wal_here
     179              :         };
     180              : 
     181              :         // If the parent is not available, what's the cheapest method involving
     182              :         // a snapshot here or later?
     183           71 :         let mut cheapest_non_incremental: Option<SegmentSize> = None;
     184           71 :         if let Some(snapshot_here) = snapshot_here {
     185           49 :             cheapest_non_incremental = Some(snapshot_here);
     186           49 :         }
     187           71 :         if let Some(snapshot_later) = snapshot_later {
     188              :             // Use <=, to prefer skipping if the size is equal
     189           39 :             if let Some(parent) = &cheapest_non_incremental {
     190           35 :                 if snapshot_later.accum_size <= parent.accum_size {
     191           34 :                     cheapest_non_incremental = Some(snapshot_later);
     192           34 :                 }
     193            4 :             } else {
     194            4 :                 cheapest_non_incremental = Some(snapshot_later);
     195            4 :             }
     196           32 :         }
     197              : 
     198              :         // And what's the cheapest method, if the parent is available?
     199           71 :         let cheapest_incremental = if let Some(cheapest_non_incremental) = &cheapest_non_incremental
     200              :         {
     201              :             // Is it cheaper to use a snapshot here or later, anyway?
     202              :             // Use <, to prefer Wal over snapshot if the cost is the same
     203           53 :             if wal_here.accum_size < cheapest_non_incremental.accum_size {
     204           32 :                 wal_here
     205              :             } else {
     206           21 :                 cheapest_non_incremental.clone()
     207              :             }
     208              :         } else {
     209           18 :             wal_here
     210              :         };
     211              : 
     212           71 :         SizeAlternatives {
     213           71 :             incremental: cheapest_incremental,
     214           71 :             non_incremental: cheapest_non_incremental,
     215           71 :         }
     216           71 :     }
     217              : }
        

Generated by: LCOV version 2.1-beta